API Reference
This page is auto-generated from in-source docstrings via mkdocstrings. When something here looks wrong, fix the docstring — not this file.
For conceptual prose (what these objects are and why they exist), see the Getting Started guide and the Concepts page.
jno.core
The top-level solver. Wraps a list of constraint expressions and a
domain, compiles them once, then exposes solve() / eval().
jno.core.core
core(constraints: List[Placeholder], mesh: Optional[Tuple[int, ...]] = (1, 1), resume_from: Optional[str] = None, *, domain: Optional[domain] = None)
core solver using traced operations.
Initialize core solver.
The random seed is read from config — JNO_SEED (env), else
[jno] seed in .jno.toml / ~/.jno/config.toml, else 42
(via jno.get_seed); it is not a constructor argument.
| PARAMETER | DESCRIPTION |
|---|---|
constraints
|
List of constraint expressions defining the problem to solve. Each constraint represents an equation or condition that should be minimized during training (e.g., PDE residuals, boundary conditions, data fitting terms).
TYPE:
|
domain
|
Optional domain override. When omitted ( A parameter-only fit needs no domain of its own: an inverse
loss like Pass
TYPE:
|
mesh
|
Shape of the device mesh for hybrid parallelism as a tuple (batch, model). Controls how computation is distributed across multiple GPUs/TPUs.
Examples:
- (1, 1): the default -- expanded to (n_devices, 1), data parallelism
over every device Note: batch * model must equal the total number of available devices;
any other shape logs a warning and falls back to (n_devices, 1). To run on
one device, make only that one visible to JAX (e.g. Recommendations: - Model fits on 1 GPU: Use (n_devices, 1) for maximum throughput - Model doesn't fit on 1 GPU: Use (1, n_devices) for model sharding - Large model + large data: Use hybrid, e.g., (2, 2) on 4 GPUs Default: (1, 1), expanded to (n_devices, 1): pure data parallelism over every visible device. Each training array is split over its samples, else over the points of a single sample, else replicated.
TYPE:
|
resume_from
|
Path to a checkpoint directory written by
:class:
TYPE:
|
solve
solve(epochs: int = 1000, batchsize: Optional[int] = None, checkpoint_gradients: bool = False, offload_data: bool = False, inner_steps: int = 1, accumulation_steps: int = 1, min_consecutive: Optional[int] = 1, profile: bool = False, callbacks: Optional[List] = None, substeps: list | None = None) -> statistics
Train using per-model optimizers attached via model.optimizer().
Every model used in the constraints must have an optimizer
attached before calling solve(). Models can optionally be
frozen (model.freeze()) or have LoRA enabled
(model.lora(rank, alpha)).
| PARAMETER | DESCRIPTION |
|---|---|
epochs
|
Number of training epochs.
TYPE:
|
batchsize
|
Mini-batch size (
TYPE:
|
checkpoint_gradients
|
If
TYPE:
|
offload_data
|
If
TYPE:
|
inner_steps
|
Number of gradient steps to fuse into a single
TYPE:
|
accumulation_steps
|
Number of micro-batches whose gradients
are averaged before a single optimizer update. The
effective batch size becomes
TYPE:
|
min_consecutive
|
Minimum number of consecutive time steps
fed to each constraint evaluation.
TYPE:
|
profile
|
If
TYPE:
|
callbacks
|
Optional list of :class:
TYPE:
|
substeps
|
Optional list of substep specs for alternating
optimisation. Each entry is either a plain list of constraint
indices Example — HyCo alternating::
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
statistics
|
Training history with
TYPE:
|
eval
eval(operation: Union[List[BinaryOp], BinaryOp], domain: Optional[domain] = None, min_consecutive: Optional[int] = 1, key: Any = None, samples: str = 'auto')
Evaluate an operation (or list of operations) on the current models.
| PARAMETER | DESCRIPTION |
|---|---|
operation
|
Expression(s) to evaluate.
TYPE:
|
domain
|
Override the stored domain.
TYPE:
|
min_consecutive
|
Consecutive-time-step window for time-dependent expressions.
TYPE:
|
key
|
Optional PRNG key for stochastic ops.
TYPE:
|
samples
|
How to handle Bayesian models in the dependency graph:
The default flips to chain automatically because a single
last-sample evaluation of a nonlinear function of Bayesian
weights is, in general, not a meaningful summary of the
posterior (
TYPE:
|
sweep
sweep(space: ArchSpace, optimizer: Union[str, type, None] = None, budget: int = 0, devices: Union[None, int, str, List[int], DeviceConfig] = None) -> statistics
Run architecture and hyperparameter search with optional parallelism.
| PARAMETER | DESCRIPTION |
|---|---|
space
|
ArchSpace defining the search space (architecture + training params)
TYPE:
|
optimizer
|
Nevergrad optimizer name (e.g., "NGOpt", "OnePlusOne", "CMA"), class, or None for exhaustive grid search
TYPE:
|
budget
|
Number of configurations to try (ignored for grid search)
TYPE:
|
devices
|
Device specification for parallel execution: - None: auto-detect and use all available devices - int: use this many devices - str: device type ("gpu", "cpu", "tpu") - List[int]: specific device indices to use - DeviceConfig: explicit device configuration
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
statistics
|
Training statistics from the best configuration |
Domain
jno.domain is the entry point for spatial geometry, mesh management,
sampling, and tensor tags.
The unknown — d.unknown
d.unknown(...) is the field you solve for, the same object for jno.fem, jno.fdm and a PINN
residual; the shape namespace spells the common cases:
u = d.unknown() # scalar, P1 — also d.unknown.scalar()
U = d.unknown.vector(2, order=2) # value_shape=(2,), P2
S = d.unknown.matrix(2, 2, symmetric=True) # 3 stored values per node, a 2x2 in every expression
E = d.unknown.vector(3, space="N1E") # an edge-element field
v = u.test() # FEM: its test function (same space, shape, order)
Keyword options: name=, order=, space=, complex= (and symmetric= on matrix), and
constant=True for ONE value (vector, matrix) over the whole domain — see
constant unknowns. In jno.fem the
unknown is the trial function and u.test() its test function; the result is bit-identical to the same
problem written with u, v = d.fem_symbols(...), which keeps working. jno.fdm solves for the values of
a real P1 nodal unknown with full storage; any other combination (order=2, a non-nodal space,
complex=True, symmetric=True) is FEM-only and jno.fdm refuses it.
domain.unknown -- callable as d.unknown(...), and the shape namespace .scalar / .vector
/ .matrix. Each builds the same kind of unknown; the namespace only spells the shape.
scalar
A scalar unknown, value_shape=(). constant=True: ONE value over the whole domain.
vector
A vector unknown with n components, value_shape=(n,). constant=True: one vector.
matrix
An n x m matrix unknown. symmetric=True (square only) stores the n(n+1)/2 upper-triangle
values per node while the symbol behaves as the full matrix in every expression. constant=True:
one matrix over the whole domain.
jno.domain.csg
Lazy Shapely-backed 2D polygon domain with true CSG operators.
The class preserves the jno.domain variable/context contract but does
not create a mesh. Point sets are materialized only when variable or
sample is called with an explicit sample count.
from_polygons
classmethod
from_polygons(polygons: Mapping[str, Sequence[Sequence[float]]], *, time: Optional[Tuple[float, float, int]] = None, compute_mesh_connectivity: bool = False, mesh_size: Optional[float] = None, sampler: Optional[Any] = None, samplers: Optional[Mapping[str, Any]] = None, resampling_strategy: Optional[Any] = None, resampling_strategies: Optional[Mapping[str, Any]] = None) -> 'PolygonDomain'
Create one CSG domain from a mapping of region names to vertices.
from_regions
classmethod
from_regions(regions: Mapping[str, BaseGeometry], *, time: Optional[Tuple[float, float, int]] = None, compute_mesh_connectivity: bool = False, mesh_size: Optional[float] = None, sampler: Optional[Any] = None, samplers: Optional[Mapping[str, Any]] = None, resampling_strategy: Optional[Any] = None, resampling_strategies: Optional[Mapping[str, Any]] = None) -> 'PolygonDomain'
Create one CSG domain from named Shapely polygonal regions.
region
Define a named sub-region addressable via domain.variable(name) and usable
as a FEM boundary-condition location.
Parameters
name:
Region tag name (e.g. "inlet", "right_top").
where:
* a shapely geometry — a boundary LineString/MultiLineString or an area
Polygon/MultiPolygon registered directly;
* a point predicate f(x, y) -> bool selecting the part of the active
boundary where it holds (evaluated per analytic edge-segment midpoint, so
it selects whole polygon edges);
* a str aliasing an already-registered tag.
kind:
"boundary" or "interior". None auto-detects from the geometry
(area → interior, line → boundary); predicates are boundary-only for now.
Returns self for chaining.
add_boundary_segments
add_boundary_segments(tag: str, segments: Sequence[Sequence[Sequence[float]]], *, normal_geometry: Optional[Any] = None) -> 'PolygonDomain'
Register an additional boundary tag from explicit line segments.
This is intended for imported boundary-condition/radiation surfaces that are subsets of component boundaries rather than whole closed polygons.
compute_enclosure_view_factor
compute_enclosure_view_factor(tags: Sequence[str], opaque_tags: Optional[Sequence[str]] = None, medium_tags: Optional[Sequence[str]] = None)
Compute cross-tag polygon boundary view factors for radiative BCs.
All tags must be polygon boundary tags that have already been sampled with normals. The method ray-traces line-of-sight against all known polygon boundary segments, then stores one visibility block and one view-factor block for every source/target tag pair:
v_<source>__<target> and f_<source>__<target>.
| PARAMETER | DESCRIPTION |
|---|---|
tags
|
Boundary tags participating in the radiation enclosure.
TYPE:
|
opaque_tags
|
Accepted for API compatibility. PolygonDomain uses all known polygon boundaries as opaque blockers, so this argument is currently informational.
TYPE:
|
medium_tags
|
Region names whose union is the radiating medium.
Normals are oriented to point into this medium before computing
view factors. If omitted and regions named
TYPE:
|
draw_candidates
Return (points, normals_or_None) candidate pool for resampling.
Generates fresh candidate points from the polygon geometry on each call (10× the currently-sampled count, min 1000) so that resampling strategies can explore the full domain rather than being confined to the initial sample.
stack
classmethod
Stack multiple PolygonDomains into one batched domain for multi-geometry training.
Use n * dom to set how many independent samplings of a geometry appear
in the training batch before passing it here. Points are drawn by rejection
sampling so n_interior / n_boundary are exact regardless of mesh size.
| PARAMETER | DESCRIPTION |
|---|---|
*batched_domains
|
PolygonDomain instances, typically
TYPE:
|
n_interior
|
Interior collocation points per geometry sample.
TYPE:
|
n_boundary
|
Boundary points per geometry sample.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
'domain'
|
A |
'domain'
|
yields |
'domain'
|
sampling per geometry instance. |
Example::
from shapely.geometry import box, Point
import jno
d1 = jno.domain(box(0, 0, 1, 1))
d2 = jno.domain(Point(0.5, 0.5).buffer(0.5))
dom = jno.domain.stack(100 * d1, 100 * d2, n_interior=512, n_boundary=128)
x, y = dom.variable("interior") # (200, 1, 512, 1) each
xb, yb = dom.variable("boundary") # (200, 1, 128, 1) each
build_mesh
build_mesh(mesh_size: float = 0.1, *, algorithm: Optional[int] = None, threads: Optional[int] = None, region_mesh_sizes: Optional[Mapping[str, float]] = None, sizes: Optional[Mapping[str, float]] = None, interpolate: bool = True) -> 'PolygonDomain'
Generate a gmsh mesh from the active Shapely CSG geometry.
After this call, self.mesh_connectivity and self._boundary_registry
are populated and downstream operations that need a mesh
(expr.integrate(), scheme="finite_difference" derivatives) become
available. The lazy sampling path is untouched: previously materialized
collocation samples in self.context survive and automatic-
differentiation derivatives keep using them.
| PARAMETER | DESCRIPTION |
|---|---|
mesh_size
|
Default target element size for points that don't fall on any per-region boundary override.
TYPE:
|
algorithm
|
gmsh 2-D meshing kernel (
TYPE:
|
threads
|
gmsh thread count (
TYPE:
|
region_mesh_sizes
|
Per-source-region mesh size overrides keyed by
the names used to construct the source regions (e.g. the
TYPE:
|
interpolate
|
Controls how
TYPE:
|
Re-calling build_mesh re-meshes from scratch and clears the integral
weight cache.
Neural-network controls
jno.nn lifts a plain Equinox / foundax module into a jNO
Model so it can participate in the trace and accept per-model
optimisers, masks, LoRA, freezing, and so on.
jno.architectures.models.nn
Neural network wrapping class for integrating modules into the jno pipeline.
Use nn.wrap(module) (or the shorthand nn(module)) to wrap an
Equinox, Flax Linen, or Flax NNX module into a Model that works with
jno.core.
Architecture factories have moved to the foundax package::
import foundax
model = jno.nn.wrap(foundax.mlp(2, hidden_dims=64, key=key))
wrap
classmethod
wrap(module: Any, space: ArchSpace = None, name: str = '', weight_path: str = None) -> Union[Model, TunableModule]
Wrap a module for use in the jno pipeline.
This is the primary method for integrating custom architectures into the jno framework. It handles both standard wrapping and architecture search scenarios.
| PARAMETER | DESCRIPTION |
|---|---|
module
|
An
TYPE:
|
space
|
Optional
TYPE:
|
name
|
Optional display name.
TYPE:
|
weight_path
|
Optional path to pretrained weights.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Model
|
Standard wrapped module (when space=None).
TYPE:
|
TunableModule
|
Tunable module for architecture search (when space provided).
TYPE:
|
| RAISES | DESCRIPTION |
|---|---|
ValueError
|
If |
Example
Wrap a custom equinox module
import foundax model = nn.wrap(foundax.mlp(2, output_dim=1, key=jax.random.PRNGKey(0)))
jno.trace.Model
Wrapper for user-defined Equinox models.
Allows using any Equinox module within the PINO tracing system. The module is initialized lazily when the input dimension is known.
Example - Direct call style (module takes separate arguments): class MLP(eqx.Module): ... def call(self, x, y, *, key=None): z = jnp.concat([x, y], axis=-1) ... return z
uv_net = pnp.nn.wrap(MLP(..., key=key))
u = uv_net(x, y)[..., 0]
Create a Model wrapper.
| PARAMETER | DESCRIPTION |
|---|---|
module
|
An Equinox module instance (already constructed), or a callable / Flax nn.Module for backward compatibility.
TYPE:
|
optimizer
Attach an optimizer to this model.
When preceded by mask(param_mask), the optimizer applies only
to matching parameters; everything else uses the global optimizer
(set via a bare optimizer() call)::
NN.mask(mask_decoder).optimizer(optax.adam) # decoder group
NN.mask(mask_encoder).optimizer(optax.sgd) # encoder group
NN.optimizer(optax.adam) # global fallback
Bake the learning rate into the optax optimizer (e.g. optax.adam(1e-3));
use :meth:scale to multiply it -- e.g. with a dlrs(...) schedule for
loss-adaptive learning-rate scaling. mask(...) is one-shot, so to scale a
masked group call mask(...) again before scale(...)::
NN.mask(mask_decoder).optimizer(optax.adam(1e-3))
NN.mask(mask_decoder).scale(my_schedule)
A bare/global call (not preceded by mask(...)) replaces any
previously configured parameter groups.
| PARAMETER | DESCRIPTION |
|---|---|
opt_fn
|
An optax optimizer factory, e.g.
TYPE:
|
freeze
Mark this model as frozen (not trained).
When preceded by mask(...), only the currently selected
parameters are frozen and everything else remains trainable::
NN.mask(param_mask).freeze() # True leaves frozen, False leaves trainable
NN.freeze() # whole model frozen
Order matters: mask() must be called before freeze().
mask
Set the current mask scope using an explicit boolean pytree mask.
param_mask must mirror the parameter tree structure and contain
boolean leaves where True selects leaves in the masked scope.
This scope is consumed by grouped optimizer/lr calls and by
mask(...).freeze(). It is also read by u.grad(net.mask(...))
to restrict the Jacobian to only the selected parameters.
Example::
import equinox as eqx, jax
all_false = jax.tree_util.tree_map(lambda _: False, model.module)
param_mask = eqx.tree_at(
lambda m: (m.layers[0].weight, m.layers[0].bias),
all_false, (True, True),
)
model.mask(param_mask).optimizer(optax.adam(1e-3))
J = crux.eval([u.grad(model.mask(param_mask))])[0] # (N, P_selected)
lora
lora(rank: int = 4, alpha: float = 1.0, *, target: str | None = None, wrapper: type[LoRAWrapper] | Sequence[type[LoRAWrapper]] | None = None, specs: list[dict] | None = None)
Enable LoRA fine-tuning for this model.
Two calling conventions:
-
Uniform::
NN.lora(rank=8, alpha=16) NN.lora(rank=4, wrapper=MyConvAdapter) # custom adapter NN.lora(rank=4, wrapper=[LoRALinear, MyConv]) # tried in order
-
Per-target — different rank/alpha/adapter per layer group::
NN.lora(specs=[ {"target": "encoder", "rank": 4, "alpha": 1.0}, {"target": "conv", "rank": 8, "alpha": 2.0, "wrapper": MyConvAdapter}, ])
Each target is a regex matched against the pytree path.
The first matching spec wins.
By default only the low-rank adapters are trained; base weights are
frozen. Layers that are NOT wrapped by LoRA remain fully trainable.
Call freeze() before lora() to also freeze any parameters
outside LoRA-wrapped layers::
NN.freeze().lora(rank=8, alpha=16)
Use mask(M) to restrict which layers receive LoRA adapters::
NN.mask(M).lora(rank=8, alpha=16) # only M-selected layers are wrapped
| PARAMETER | DESCRIPTION |
|---|---|
rank
|
LoRA rank (uniform mode).
TYPE:
|
alpha
|
LoRA scaling factor (uniform mode).
TYPE:
|
target
|
Regex to restrict which layers get LoRA adapters (uniform
mode only). Layers whose pytree path does not match are left
completely untouched. Use
TYPE:
|
wrapper
|
Adapter class or list of classes to try in order.
Defaults to
TYPE:
|
specs
|
Per-target specs (per-target mode). Each dict has keys
TYPE:
|
dtype
Set this model's working dtype (parameters and compute).
Casts all floating-point parameters to dtype and — at the forward
seam — casts the model's inputs to match, so the network actually
computes in dtype rather than promoting back to float32. The cast is
symmetric: it lowers (float32 → bfloat16) and promotes (load a
bfloat16 checkpoint, then .dtype(jnp.float32)), and applies to both
training and inference. Integer arrays (e.g. indices) are left unchanged.
This is the model-precision knob. Data precision (float32 vs
float64) is JAX's jax_enable_x64 flag — not a jNO setting. Enable
it before building models/domains (JAX_ENABLE_X64=1 or
jax.config.update("jax_enable_x64", True)).
| PARAMETER | DESCRIPTION |
|---|---|
dtype
|
A JAX floating dtype object, e.g.
TYPE:
|
Caveats
- bfloat16 compute degrades autodiff derivatives
(
.laplacian/.hessian) — keep derivative-critical (PINN) models in float32 and opt only data-loss / operator backbones into bf16. - bfloat16 parameters mean the optimizer update also runs in bfloat16, which can stall on very small updates.
Example::
backbone.dtype(jnp.bfloat16) # real bf16 compute for this model
pinn_net.dtype(jnp.float32) # keep its derivatives full precision
constrain
Apply a paramax reparameterization to trainable parameter leaves.
Parameters are stored in their unconstrained form and transformed by
transform before every forward pass via paramax.unwrap(),
which jno's training loop calls automatically.
When preceded by mask(...), only leaves where the mask is True
are wrapped — all other leaves remain unconstrained::
k_net.mask(output_mask).constrain(jax.nn.softplus) # output layer only
k_net.constrain(jax.nn.softplus) # all parameters
| PARAMETER | DESCRIPTION |
|---|---|
transform
|
A jit-compatible callable (e.g.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
'Model'
|
self (for chaining) |
initialize
Load pretrained weights into this model at init time.
Accepted ``weights`` inputs:
- ``str`` / ``Path``: load from checkpoint path.
Supports Equinox ``.eqx`` files and Orbax checkpoint directories
(optionally ``"<path>::<model_key>"``).
- Pytree object: copy array leaves directly from the provided tree.
- Callable initializer: apply a JAX initializer function to every
floating-point array leaf at compile time.
Examples:
.. code-block:: python
net.initialize("./weights.eqx")
net.initialize("./runs/ckpts/2000::1")
net.initialize(other_model.module)
p = jno.np.parameter((1,), key=jax.random.PRNGKey(0))
p.initialize(jax.nn.initializers.ones)
| PARAMETER | DESCRIPTION |
|---|---|
weights
|
File path / pytree / callable initializer.
TYPE:
|
key
|
Optional PRNG key used when
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
'Model'
|
self (for chaining). |
tune
tune(*, freeze: list | None = None, lora: list | None = None, optimizer: list | None = None, lr: list | None = None, dtype: list | None = None) -> 'Model'
Declare per-model tunable options for hyperparameter sweeps.
Each argument accepts a list of candidate values. During a sweep the tuner searches over all combinations.
| PARAMETER | DESCRIPTION |
|---|---|
freeze
|
List of bool, e.g.
TYPE:
|
lora
|
List of
TYPE:
|
optimizer
|
List of optax factories, e.g.
TYPE:
|
lr
|
List of :class:
TYPE:
|
dtype
|
List of dtypes, e.g.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
'Model'
|
self (for chaining). |
Example::
backbone = nn.poseidon(...)
backbone.initialize("weights.msgpack")
backbone.tune(
freeze=[True, False],
lora=[(4, 1.0), None],
optimizer=[optax.adam],
lr=[lrs.constant(1e-4), lrs.constant(1e-5)],
)
Symbolic math (jno.np)
A NumPy-compatible namespace that returns traced placeholders instead
of concrete arrays. Use it inside any expression that you intend to
feed into jno.core(...).
jno.jnp_ops
concat
Concatenate placeholders along an axis (always axis=-1 at eval time).
grad
grad(target: Placeholder, variable: Variable, scheme: str = 'automatic_differentiation') -> Jacobian
Compute the gradient of target with respect to variable.
Implemented as a single-variable Jacobian.
Prefer the method-style shorthand on the target expression::
u_x = u.d(x) # ∂u/∂x
u_xx = u.d(x).d(x) # ∂²u/∂x² (chainable)
| PARAMETER | DESCRIPTION |
|---|---|
target
|
Expression to differentiate
TYPE:
|
variable
|
Variable to differentiate with respect to
TYPE:
|
scheme
|
'automatic_differentiation' (default) or 'finite_difference'
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Jacobian
|
Jacobian placeholder representing ∂target/∂variable |
Example
u_x = pnp.grad(u(x, y), x) # ∂u/∂x
Function helpers and loss balancers (jno.fn)
jno.fn provides PDE-named helpers (heat, wave, burgers_1d, ...),
loss reductions (mse, mae, rmse, huber, log_cosh, ...), and
the adaptive loss balancers under jno.fn.adaptive.*.
jno.fn
Functional helpers for traced expressions: jno.fn.sin(u), jno.fn.mse(pred, target).
This module is callable — jno.fn(my_func, [arg1, arg2]) wraps an
arbitrary function into the tracing graph (replaces jno.np.function).
Sections
- Math: sin, cos, exp, log, sqrt, abs, …
- Losses: mse, mae, rmse, huber, log_cosh, relative_l2
- PDEs: poisson, heat, wave, burgers_1d, navier_stokes_incompressible_2d, …
Examples
import jno pde = jno.fn.sin(u) + jno.fn.exp(-x) loss = jno.fn.mse(pred, target) custom = jno.fn(lambda a, b: a ** 2 + b, [u, v])
_module_call
_module_call(fn: Callable, args: list = [], name: str = '', reduces_axis: Optional[int] = None) -> FunctionCall
Wrap an arbitrary function into the tracing graph.
| PARAMETER | DESCRIPTION |
|---|---|
fn
|
Any callable
TYPE:
|
args
|
Traced placeholder arguments.
TYPE:
|
name
|
Optional display name in the expression tree.
TYPE:
|
reduces_axis
|
If the function reduces an axis, specify it here.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
FunctionCall
|
|
Example::
custom = jno.fn(lambda a, b: a ** 2 + b, [u, v], name="my_op")
Training history
solve() returns a statistics object. The most common operations:
jno.utils.statistics.statistics
Training history returned by core.solve().
Access patterns::
history = crux.solve(...)
history.total_loss # final scalar total loss
history.total_loss_history # 1-D array of total loss per epoch
history.training_logs # list of per-solve()-call dicts
history.training_logs[-1]["total_loss"] # full array from last call
history.plot("./runs/loss.png") # quick visualization
total_loss
property
Final scalar total loss (last value across all solve() calls).
Returns None when no training has been recorded.
total_loss_history
property
1-D array of total loss concatenated across all solve() calls.
plot
Plot training statistics from all solve() calls.
Creates a multi-panel figure showing: - Constraint losses over time (individual lines; total added when >1 constraint) - Tracker values over time (if any trackers were defined) - Step time in milliseconds (derived from log timestamps)
| PARAMETER | DESCRIPTION |
|---|---|
path
|
Path to save the figure (e.g. "./runs/training.png").
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
statistics
|
self (for chaining) |
load
classmethod
Load a trained core model from a file.
Restores all trained parameters, operations, domain, and history.
| PARAMETER | DESCRIPTION |
|---|---|
filepath
|
Path to saved model file
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
statistics
|
core instance with trained parameters |
Example
sol = core.load("trained_model.pkl")
Differential and integral operators
These provide the residuals you put inside constraints
(u.laplacian(x, y), u.d(x), (grad_u * n).integrate()).
Scheme strings
Every differential operator (.d, .diff, .d2, .dd, .laplacian,
.hessian) accepts a scheme= kwarg that selects the backend:
| Scheme | Backend |
|---|---|
"automatic_differentiation" (default) |
global default — see jno.setup(diff_type=..., hessian_type=...) |
"automatic_differentiation:forward" |
first-order via jax.jacfwd |
"automatic_differentiation:reverse" |
first-order via jax.jacrev |
"automatic_differentiation:fwd-over-rev" |
second-order jacfwd(jacrev(f)) (= historical jax.hessian) |
"automatic_differentiation:fwd-over-fwd" |
second-order jacfwd(jacfwd(f)) |
"automatic_differentiation:rev-over-rev" |
second-order jacrev(jacrev(f)) |
"automatic_differentiation:rev-over-fwd" |
second-order jacrev(jacfwd(f)) |
"finite_difference" |
central-difference stencils on mesh (with :lsq / :uniform / :inverse_distance / :cotangent sub-schemes) |
"spectral" |
FFT along the grid axes — exact for band-limited periodic fields; uniform grid only |
"spectral:cosine" |
even (mirror) extension — for fields with vanishing odd derivatives at both ends |
jno.fd(order=4), jno.fd(points=…), jno.fd(weights=…), boundary=, upwind=, average=, fit=, rings= |
any finite-difference stencil: on a structured grid by its order or its points (Fornberg weights), on an unstructured mesh by a local polynomial fit; also u.bind(x=x, y=y, scheme=jno.fd(...)) for every .x/.xx of that binding |
Forward-mode is typically cheaper when the input dim (≤ 3 spatial dims for
PINNs) is ≤ the output dim; reverse-mode is cheaper for scalar losses with
many inputs. Set the project-wide default once via .jno.toml:
[jno]
diff_type = "forward" # default for first-order operators
hessian_type = "fwd-over-rev" # default for second-order operators
or per script via jno.setup(__file__, diff_type="forward"). Per-call
scheme= always overrides the default.
diff_type accepts a whole scheme as well as an AD sub-mode, so
jno.setup(__file__, diff_type="spectral") makes every unqualified derivative in the run spectral
— see Operations → Spectral differentiation for the
accuracy numbers and the periodicity caveat.
jno.differential_operators.DifferentialOperators
Static collection of mesh-based FD operators (1-D, 2-D, 3-D).
All public methods are static — the class is used purely as a namespace. See the module docstring for full method descriptions.
compute_fd_gradient_1d_simple
staticmethod
compute_fd_gradient_1d_simple(u_values: ndarray, points: ndarray, lines: ndarray, method: str = 'area_weighted') -> jnp.ndarray
Gradient on a 1-D line mesh.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values at mesh points, shape
TYPE:
|
points
|
Mesh point coordinates, shape
TYPE:
|
lines
|
Line element connectivity, shape
TYPE:
|
method
|
Weighting strategy — one of
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
|
compute_fd_laplacian_1d_simple
staticmethod
compute_fd_laplacian_1d_simple(u_values: ndarray, points: ndarray, lines: ndarray, method: str = 'gradient_of_gradient') -> jnp.ndarray
Laplacian on a 1-D line mesh.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
lines
|
Line connectivity, shape
TYPE:
|
method
|
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
|
compute_fd_hessian_1d_simple
staticmethod
compute_fd_hessian_1d_simple(u_values: ndarray, points: ndarray, lines: ndarray, var_dims: list | None = None) -> jnp.ndarray
Hessian (= d²u/dx²) on a 1-D line mesh.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
lines
|
Line connectivity, shape
TYPE:
|
var_dims
|
Optional
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
Hessian, shape |
compute_fd_gradient_2d_simple
staticmethod
compute_fd_gradient_2d_simple(u_values: ndarray, points: ndarray, triangles: ndarray, dim: int, method: str = 'area_weighted', grid: dict | None = None) -> jnp.ndarray
Gradient on a 2-D triangular mesh.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values at mesh points, shape
TYPE:
|
points
|
Mesh point coordinates, shape
TYPE:
|
triangles
|
Triangle connectivity, shape
TYPE:
|
dim
|
Spatial dimension to differentiate (0 = x, 1 = y).
TYPE:
|
method
|
TYPE:
|
grid
|
Optional structured-grid descriptor
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
|
compute_gradient_2d_lsq
staticmethod
compute_gradient_2d_lsq(u_values: ndarray, points: ndarray, triangles: ndarray, dim: int) -> jnp.ndarray
Least-squares gradient on a 2-D triangular mesh.
For each node i the gradient is estimated by solving a 2×2 area-weighted least-squares problem built from incident triangle centroids. The 2×2 system is solved via Cramer's rule so the entire computation uses only JAX scatter-add operations with no per-node Python loops.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
triangles
|
Triangle connectivity, shape
TYPE:
|
dim
|
0 →
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
Gradient component at each node, shape |
compute_fd_laplacian_2d_simple
staticmethod
compute_fd_laplacian_2d_simple(u_values: ndarray, points: ndarray, triangles: ndarray, dims: tuple, method: str = 'gradient_of_gradient', grid: dict | None = None) -> jnp.ndarray
Laplacian on a 2-D triangular mesh.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
triangles
|
Triangle connectivity, shape
TYPE:
|
dims
|
Spatial dimensions to sum over, e.g.
TYPE:
|
method
|
TYPE:
|
grid
|
Optional structured-grid descriptor
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
Laplacian, shape |
compute_laplacian_2d_cotangent
staticmethod
compute_laplacian_2d_cotangent(u_values: ndarray, points: ndarray, triangles: ndarray) -> jnp.ndarray
Cotangent-weight (Laplace–Beltrami) Laplacian on a 2-D mesh.
For each triangle (i, j, k) the cotangent of each interior angle
is used to weight the edge contributions::
lap[i] += (1/A_i) * [ cot_k*(u_j - u_i) + cot_j*(u_k - u_i) ]
with cot_k = cotangent of the angle at vertex k (opposite edge
(i,j)), and A_i = (1/3) * Σ area.
This is second-order accurate and isotropic; it is the gold standard for PDE discretisation on unstructured 2-D triangular meshes.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
triangles
|
Triangle connectivity, shape
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
Laplacian at each point, shape |
compute_fd_hessian_2d_simple
staticmethod
compute_fd_hessian_2d_simple(u_values: ndarray, points: ndarray, triangles: ndarray, var_dims: list, grid: dict | None = None, method=None) -> jnp.ndarray
Hessian on a 2-D triangular mesh (area-weighted FD).
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
triangles
|
Triangle connectivity, shape
TYPE:
|
var_dims
|
List of
TYPE:
|
grid
|
Optional structured-grid descriptor (from
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
Hessian, shape |
compute_fd_gradient_3d_simple
staticmethod
compute_fd_gradient_3d_simple(u_values: ndarray, points: ndarray, tetrahedra: ndarray, dim: int, method: str = 'area_weighted', grid: dict | None = None) -> jnp.ndarray
Gradient on a 3-D tetrahedral mesh.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Mesh point coordinates, shape
TYPE:
|
tetrahedra
|
Tet connectivity, shape
TYPE:
|
dim
|
Spatial dimension (0, 1 or 2).
TYPE:
|
method
|
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
|
ndarray
|
the reshaped central difference is used (see the 2-D gradient); |
compute_gradient_3d_lsq
staticmethod
compute_gradient_3d_lsq(u_values: ndarray, points: ndarray, tetrahedra: ndarray, dim: int) -> jnp.ndarray
Least-squares gradient on a 3-D tetrahedral mesh.
Analogous to :meth:compute_gradient_2d_lsq but solves a 3×3
normal-equation system at each node via Cramer's rule.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
tetrahedra
|
Tet connectivity, shape
TYPE:
|
dim
|
0 →
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
Gradient component at each node, shape |
compute_laplacian_3d_cotangent
staticmethod
compute_laplacian_3d_cotangent(u_values: ndarray, points: ndarray, tetrahedra: ndarray) -> jnp.ndarray
P1 finite-element (Laplace–Beltrami) Laplacian on a tetrahedral mesh — the 3-D analogue of the
2-D :meth:compute_laplacian_2d_cotangent (in 2-D the cotangent weights are the P1 stiffness
off-diagonals).
Each tet's linear basis gradients ∇φ_a are constant, obtained from the inverse of the element
Jacobian J = [p1-p0 | p2-p0 | p3-p0] (∇λ_{1,2,3} are the rows of J⁻¹, ∇λ_0 their
negated sum). The assembled stiffness applied to u is, matrix-free,
(K u)_i = Σ_{T∋i} V_T (∇φ_i · ∇u_T) with ∇u_T = Σ_a u_a ∇φ_a the (constant) element
gradient and V_T the tet volume; normalizing by the lumped vertex volume
M_i = Σ_{T∋i} V_T / 4 gives the strong Laplacian Δu_i = -(K u)_i / M_i. Symmetric,
CG-compatible, and second-order for the Galerkin solution — the accurate default on tets.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
tetrahedra
|
Tet connectivity, shape
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
Laplacian at each node, shape |
compute_fd_laplacian_3d_simple
staticmethod
compute_fd_laplacian_3d_simple(u_values: ndarray, points: ndarray, tetrahedra: ndarray, dims: tuple, method: str = 'gradient_of_gradient', grid: dict | None = None) -> jnp.ndarray
Laplacian on a 3-D tetrahedral mesh.
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
tetrahedra
|
Tet connectivity, shape
TYPE:
|
dims
|
Spatial dimensions to sum over, e.g.
TYPE:
|
method
|
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
Laplacian, shape |
compute_fd_hessian_3d_simple
staticmethod
compute_fd_hessian_3d_simple(u_values: ndarray, points: ndarray, tetrahedra: ndarray, var_dims: list, grid: dict | None = None, method=None) -> jnp.ndarray
Hessian on a 3-D tetrahedral mesh (volume-weighted FD).
| PARAMETER | DESCRIPTION |
|---|---|
u_values
|
Function values, shape
TYPE:
|
points
|
Coordinates, shape
TYPE:
|
tetrahedra
|
Tet connectivity, shape
TYPE:
|
var_dims
|
List of
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
Hessian, shape |
parse_fd_scheme
staticmethod
Parse a scheme string into (main_scheme, grad_method, lap_method).
Supported formats::
"finite_difference" → fd, "area_weighted", "gradient_of_gradient"
"finite_difference:lsq" → fd, "least_squares", "lsq_of_gradient"
"finite_difference:cotangent" → fd, "area_weighted", "cotangent"
"finite_difference:uniform" → fd, "uniform", "gradient_of_gradient"
"finite_difference:inverse_distance" → fd, "inverse_distance", "gradient_of_gradient"
"automatic_differentiation" → ad, None, None
| RETURNS | DESCRIPTION |
|---|---|
tuple[str, str, str]
|
Tuple |
jno.integration_operators.IntegrationOperators
Static namespace for mesh-based numerical integration.
Works on the mesh_connectivity dict produced by the domain class.
Boundary weights (nodal_ds) are already stored there; this class
adds volume weights (nodal_volumes) computed on the fly.
nodal_volumes
staticmethod
Per-node volume weights for interior integration.
Returns mesh_connectivity["nodal_volumes"] if it was precomputed
during domain setup (the normal path). Otherwise computes on the fly
(fallback for manually constructed mesh_connectivity dicts).
Each node receives a share of surrounding element volumes:
- 1-D: ½ × sum of adjacent segment lengths (trapezoidal rule)
- 2-D: ⅓ × sum of incident triangle areas
- 3-D: ¼ × sum of incident tetrahedron volumes
Parameters
mesh_connectivity : dict Preprocessed mesh connectivity from the domain class.
Returns
vols : ndarray of shape (n_points,)
gauss_points_and_weights
staticmethod
Element Gauss quadrature over the volume mesh: physical points + JxW weights.
Unlike :meth:nodal_volumes (a vertex rule that samples only at mesh nodes), this maps a
reference-cell Gauss rule of the requested degree into every element and returns the
physical quadrature points together with their w · |det J| weights, so
sum(f(points) * weights) is the higher-order Gauss approximation of ∫ f dx. Exact
for polynomials up to degree; many points per element make it far harder for an
expressive integrand (e.g. a network) to alias the rule than the vertex rule can.
Parameters
mesh_connectivity : dict
Preprocessed mesh connectivity from the domain class (points + triangles /
tetrahedra).
degree : int
Polynomial degree the rule integrates exactly (basix quadrature degree).
cells : ndarray, optional
Element connectivity subset (e.g. a sub-region's triangles); defaults to all cells.
Returns
(points, JxW) : tuple[ndarray, ndarray]
points of shape (n_cell * n_qp, dim) and JxW of shape (n_cell * n_qp,).
jno.utils.ad_mode
AD mode (forward / reverse) selection for jNO operators.
Two layers, in order of precedence:
-
Per-call scheme suffix on the operator:
-
First-order (
.d,.diff,d/dt)::u.d(x, scheme="automatic_differentiation:forward") u.d(x, scheme="automatic_differentiation:reverse")
-
Second-order (
.laplacian,.hessian,.d2,.dd)::u.laplacian(x, y, scheme="automatic_differentiation:fwd-over-rev") u.laplacian(x, y, scheme="automatic_differentiation:fwd-over-fwd") u.laplacian(x, y, scheme="automatic_differentiation:rev-over-rev") u.laplacian(x, y, scheme="automatic_differentiation:rev-over-fwd")
-
Global default — set via :func:
jno.setupor via.jno.toml::jno.setup(file, diff_type="forward", hessian_type="fwd-over-fwd")
.. code-block:: toml
[jno]
diff_type = "forward" # first-order default
hessian_type = "fwd-over-rev" # second-order default
The plain string "automatic_differentiation" (no suffix) resolves to the
current global default. Defaults match historical behaviour: first-order
reverse (was jax.jacobian = jacrev); second-order fwd-over-rev
(was jax.hessian = jacfwd ∘ jacrev).
parse_ad_scheme
Resolve a first-order scheme string to "forward" or "reverse".
Supported::
"automatic_differentiation" → global default (get_ad_mode())
"automatic_differentiation:forward" → "forward"
"automatic_differentiation:reverse" → "reverse"
parse_hessian_scheme
Resolve a second-order scheme string to (outer, inner) AD modes.
The result composes as outer(inner(f)). E.g. ("forward", "reverse")
means jax.jacfwd(jax.jacrev(f)) — the historical jax.hessian path.
Supported::
"automatic_differentiation" → global default (get_hessian_mode())
"automatic_differentiation:fwd-over-rev" → ("forward", "reverse")
"automatic_differentiation:fwd-over-fwd" → ("forward", "forward")
"automatic_differentiation:rev-over-rev" → ("reverse", "reverse")
"automatic_differentiation:rev-over-fwd" → ("reverse", "forward")
First-order suffixes forward/reverse are accepted as shorthand for
the matching same-mode composition (forward → fwd-over-fwd).
rowwise_jacobian
(len(rows), P) reverse-mode Jacobian of f: x -> (N,), with no vmap.
jax.jacrev takes one vjp and then vmaps the pullback across the rows of the
identity basis. A differentiable FEM solve bottoms out in
jax.experimental.sparse.linalg.spsolve, which JAX ships without a batching rule; jNO
registers one (:mod:jno.utils.solver.sparse_batching), so jacrev and jacfwd through
fem.solve() work -- but the batched pullback solves B adjoint systems as one
block-diagonal system, factorising the matrix once PER ROW and holding the B copies at once.
This does what jacrev does minus the vmap: ONE jax.vjp, so the forward pass and
its residuals are computed once and shared, followed by one pullback per requested row --
same factorisation count, one copy of the matrix at a time. The rows are unrolled at trace
time.
rows is an explicit sequence of output indices, so a caller that needs only some rows
-- MMA needs the inequality-constraint rows and not the objective -- pays for those only.
| PARAMETER | DESCRIPTION |
|---|---|
f
|
Callable from the pytree
TYPE:
|
x
|
The pytree to differentiate with respect to.
TYPE:
|
rows
|
Output indices to build rows for, in the returned order.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Array
|
|
Array
|
flattened in |
Array
|
produces once its per-leaf blocks are raveled and concatenated. |
Solvers and preconditioners
jno.solve and jno.precond are the slots that fem.solve(linear=…, nonlinear=…,
precond=…, time=…) composes (see the FEM guide). The families:
| Kind | jno.solve |
|---|---|
| Linear — direct | lu (sparse LU), dense |
| Linear — iterative (Krylov) | cg, bicgstab, gmres, fgmres, minres, cocg (complex-symmetric); lstsq (LSQR, least-squares); chebyshev (polynomial) |
| Linear — multigrid | amg (GPU AMG / NVIDIA AmgX via jaxamg) |
| Nonlinear | newton, picard, staggered — picard/staggered take anderson=m (Anderson acceleration) |
| Eigenproblem | eigs (generalized Kx = λMx) — dense reduction, preconditioned LOBPCG with precond=, or interior modes nearest a shift with sigma= |
| Singular values | svd (partial SVD of a rectangular, matrix-free operator — POD bases, inverse-problem ill-posedness) |
| Matrix functions (stochastic Lanczos, matrix-free) | logdet, trace, applyfun (f(A)·v), diagonal |
| Time integration | theta (θ-method), bdf2 (2nd-order and L-stable), sdirk (L-stable SDIRK, order 2/3), rosenbrock (linearly implicit, order 3/2 — no Newton), exponential (exponential integrator), adaptive (step-doubling adaptive step size) |
Matrix functions — what Ax = b cannot express
logdet, trace, applyfun and diagonal touch the operator only through its matvec, so they scale
where a factorization cannot, and they are differentiable. They answer questions a linear solve
cannot: the Bayesian log-evidence of a FEM precision, an effective-degrees-of-freedom count, a
per-DOF uncertainty map, and one exact exponential-integrator step.
A, _ = jno.fem([ui * vi + ui.x * vi.x + ui.y * vi.y - 1.0 * vi, u(xb, yb) - 0.0]).operator
jno.solve.logdet(A, samples=64) # log det A — log-evidence
jno.solve.trace(A, fun=lambda z: 1 / z, samples=64) # tr(A⁻¹) — effective DOFs
jno.solve.diagonal(A, fun=lambda z: 1 / z, samples=256) # diag(A⁻¹) as a FIELD, plottable on the mesh
jno.solve.applyfun(A, u0, fun=lambda z: jnp.exp(-dt * z)) # exp(-dt·A)·u₀ — one exact step
Accuracy on a 198-DOF FEM precision operator (cond 45)
| quantity | estimate | exact | rel |
|---|---|---|---|
logdet (samples=64) |
168.99 | 167.22 | 1.1e-02 |
trace(1/z) (samples=64) |
119.65 | 118.20 | 1.2e-02 |
diagonal(1/z) (samples=256) |
— | — | 1.1e-01 (L2 over the field) |
applyfun exp step |
— | — | 1.1e-15 |
The first three are stochastic — Hutchinson probes plus Lanczos quadrature — so a percent or
so is the expected accuracy, not a defect: variance falls with samples, bias with order.
applyfun is deterministic (a Krylov approximation, no probes) and essentially exact.
Differentiating works through the estimator: d(log det cA)/dc came back 152.198 against the
closed form n/c = 152.308.
order must stay below the Krylov dimension — and a pinned FEM operator's is small
Lanczos can only build a subspace as large as the number of distinct eigenvalues the probe sees. A jNO FEM operator has far fewer than it has rows: every Dirichlet-pinned DOF is an identity row, so eigenvalue 1.0 carries the pinned count as its multiplicity.
Measured on 2-D Poisson at mesh 0.25 — n=30, 16 pinned rows, only 15 distinct eigenvalues —
the default order=25 overran that and logdet returned NaN. It now raises instead, naming the
cause. The operator above is fine (198 rows, 151 distinct); a coarse mesh with a large boundary
fraction is not. Lower order, or apply the estimator to the free-DOF operator.
applyfun is not affected: its order is an upper bound, and it stops at the order that has
actually converged (see below). The stochastic three have no such ladder, so order is a real
request there.
applyfun picks its own order
order is an upper bound on the Krylov dimension, not an exact request — the same meaning
maxiter has for every iterative solver here. Running past the dimension a problem supports used
to degrade catastrophically and silently, because the Lanczos sub-diagonal does not collapse to
zero as a textbook "happy breakdown" would; it explodes (0.27 → 2.2 → … → 184.7), the
residual being pure round-off whose normalisation gives basis vectors of noise. On the 30-DOF
operator above, exp(A)·1 against a true 49.02:
order |
before | now |
|---|---|---|
| 15 | 3.35e-15 | 3.26e-15 |
| 20 | 1.44e-10 | 3.26e-15 |
| 25 | 5.11e+35 | 3.17e-15 |
| 29 | 1.85e+74 | 3.17e-15 |
Note order 20 — this is not only a fix for the catastrophic end, it is more accurate wherever round-off has begun to contaminate the basis. The rule is the standard a-posteriori one (Saad, SIAM J. Numer. Anal. 29(1), 1992, §4): accept the first order whose approximation agrees with its predecessor. Every nested approximation comes from the same decomposition, so it costs small dense eigendecompositions and no extra matvecs — measured overhead 0.17 ms at n=513 and 0.41 ms at n=8355.
Eigenproblems at scale
jno.solve.eigs / FEM.eigs have three paths, chosen by the arguments. With none of the iterative
arguments the pencil is reduced densely — exact, and right when you want the whole low spectrum of
a small problem, but it materializes the operator (O(N²) memory). Passing precond= selects
preconditioned LOBPCG (Knyazev, SIAM J. Sci. Comput. 23(2), 517–541, 2001), which only applies
K/M as matvecs and so runs where the dense reduction cannot. Passing sigma= targets the k
eigenvalues nearest the shift — interior modes (a cavity resonance inside a band, a Brillouin-zone
point away from the band edge), which no extremal-end iteration can reach — by shift-invert with a
depth-2 block Krylov space per sweep (Ericsson & Ruhe, Math. Comp. 35, 1980; Golub & Underwood,
The block Lanczos method, 1977): θ = 1/(λ−σ) makes the near-σ modes dominant with enormous
transformed gaps, so the transformation is its own preconditioner and precond= is rejected there. A
wide degenerate cluster at the same distance as the wanted modes — the hundreds of exact zeros of a
curl-curl (Maxwell) operator, with σ between 0 and the first cavity mode — converges too: the restart
keeps at most k copies of any one Ritz value, so the cluster cannot crowd out a wanted mode. The inner solves against K − σM default to a
host sparse LU factorized once (every sweep is then triangular substitutions); linear= swaps in a
different inner solver when a factorization is too big. Constrained pencils (Dirichlet pins, periodic
ties) compose: the reduced K − σM is assembled sparsely through the same triplet remap the periodic
solve reduction uses.
K is the source-less jno.fem whose bilinear form is the stiffness; mass= takes the mass form
as a plain term list, which eigs assembles onto the same space for you:
u, v = d.fem_symbols()
xi, yi, _ = d.variable("interior", split=True)
ui, vi = u.bind(x=xi, y=yi), v.bind(x=xi, y=yi)
K = jno.fem([ui.x * vi.x + ui.y * vi.y]) # stiffness — no source term
lam, X = K.eigs(mass=[ui * vi], k=6) # dense; λ = ω² on a Neumann box
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
jno.solve.eigs(...) is the lower-level form of the same thing and takes two assembled operators
rather than a fem and a term list — jno.solve.eigs(k=6)(K.operator[0], M).
Sweeps warm-start. X0= seeds the LOBPCG block with eigenvector guesses in the full DOF space
(restricted through any constraint elimination for you) — the classic sweep accelerator: a
parameter/frequency/k-point sweep passes each point the previous point's eigenvectors, so the
iteration only tracks the drift instead of re-finding the subspace from random. Fewer columns than
k are padded with the seeded random block. X0= is rejected on the dense path (it would be
silently ignored) and not yet wired into sigma= (the transformation converges from random in a
handful of sweeps anyway):
lam, X = K.eigs(mass=mass, k=6, precond=jno.precond.jacobi()) # first sweep point
lam2, X2 = K2.eigs(mass=mass, k=6, precond=jno.precond.jacobi(), X0=X) # next point: warm
The Rayleigh–Ritz runs in the M-inner product, so an ordinary FEM form's consistent (non-lumped)
mass matrix is handled directly, and XᵀMX = I holds on both paths. Eigenvalues are differentiable on
both — for simple eigenvalues; a degenerate cluster makes ∂λ/∂θ ill-defined either way (use the
trace of the cluster). LOBPCG freezes the converged eigenvector and differentiates the Rayleigh
quotient, which gives that derivative exactly without differentiating through the sweeps, but its
eigenvectors carry no gradient where the dense path's do.
tol/maxiter tune the iterative paths and 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: on
an ill-conditioned pencil the residual floors well above it (≈4.4e-8 on a singular all-Neumann
Laplacian with cond(K) ≈ 2e16), and a tolerance below that floor burns the budget and
NaN-poisons the LOBPCG result — which is the deliberate contract for an exhausted budget, never a
quietly under-converged spectrum. The shift-invert gate measures the original pencil's residual of
the k returned pairs (a θ-space gate would flatter it), and a shift landing exactly ON an eigenvalue
makes K − σM singular — the garbage its factorization produces fails the same gate; perturb σ off
the eigenvalue. A shift-invert call that fails its gate raises ShiftInvertNotConverged (in
jno.utils.solver.eigen, a RuntimeError) naming the cause — exhausted budget or singular shift —
rather than returning NaN (also under a plain jax.grad); only under jit/vmap, where nothing can
be raised, is its result NaN-poisoned instead.
Singular values — jno.solve.svd
eigs solves the symmetric pencil Kx = λMx. The two questions that are not eigenproblems need
the SVD of a possibly rectangular map, via Golub–Kahan bidiagonalization (Golub & Kahan,
J. SIAM Numer. Anal. Ser. B 2(2), 1965):
U, s, Vt = jno.solve.svd(snapshots, k=6) # POD basis from a (n_time, n_dofs) trajectory
U, s, Vt = jno.solve.svd(jacobian_op, k=20) # ill-posedness of a parameter-to-observable map
- POD / reduced-order models — the singular vectors are the energy-optimal basis and
ssays how many modes the trajectory actually needs. - Ill-posedness — the singular spectrum of the parameter-to-observable map says which parameter modes are recoverable at all; those below the noise floor are not, whatever the optimizer does.
A is touched only through its matvec, so it can be the JVP of a differentiable FEM solve rather than
an assembled matrix, and s differentiates back to whatever that matvec closes over.
depth (bidiagonalization steps, default 2k+10) must exceed k — the Ritz values converge from
below, so at depth == k only the largest singular value is meaningful (measured 95 % error on the
rest, against ~1e-15 at depth = 2k). Convergence is fast on the decaying spectra that make POD and
ill-posedness analysis worth doing, and slow on clustered ones (~3 % error at depth = 4k on a tight
cluster) — inspect s for a plateau if the spectrum may be flat.
What runs compiled
A slot-composed solve runs as one compiled program where it can, rather than calling the Krylov
iteration from eager Python and paying dispatch on every step. How much that is worth depends on the
device: the eager cost is host-bound and so barely varies between machines, while the compiled cost
is device-bound — the faster the GPU, the larger the ratio. On an RTX 3070 at 13759 DOFs,
bicgstab + jacobi 114.1 ms → 18.1 (6.3x), cg + jacobi 97.5 → 14.5, minres + jacobi 115.2 →
20.2, gmres + jacobi 398.4 → 183.2, fgmres + jacobi 536.3 → 300.8 (1.8x — jNO's own restart loop
does more real arithmetic, so less of its time was dispatch). On CPU, 1.8–4.2x; on a faster GPU,
bicgstab + jacobi reached 16.6x. Same answers throughout. Nothing to switch on; write the
slots as usual, and
write them inline if you like (fem.solve(linear=jno.solve.cg(), precond=jno.precond.jacobi())) —
equivalently configured specs share one compilation, so a solve in a loop compiles once.
These combinations stay eager, and are correct but not accelerated:
| Slot | Why |
|---|---|
solve.chebyshev, precond.chebyshev |
measures spectrum bounds, then branches on what it measured |
precond.amg unbuilt, precond.ams, precond.form |
assembles an auxiliary operator host-side (scipy / pyamg) |
precond.jaxamg, solve.amg |
AmgX builds its hierarchy from the matrix values; unverified under a tracer |
solve.lu, solve.dense, solve.amg |
one direct call — no per-iteration dispatch to remove |
| a bare callable in either slot | jNO knows nothing about it, so it makes no assumption |
a multi-device (shard=) solve |
already compiles itself, with the operator partitioned |
First-run compilation cost
Compiling is not free the first time. Building a 13.8k-DOF 2-D Poisson problem issues ~209 XLA
compilations totalling ~3.7 s — assembly evaluates reference-element expressions whose every distinct
operation and shape is its own small program. The cost is paid once per distinct mesh shape, and
it is cached within a process: rebuilding the same problem costs ~320 ms, and rebuilding with freshly
constructed term objects costs the same (the cache is keyed on shapes, not on object identity). A
different mesh pays it again, so a remeshing loop (fem.adapt, a refinement study) pays it per
iteration.
Across separate processes — running a script twice, a test suite, CI — nothing is reused by default. JAX can persist compilations to disk, which takes the same build from 4757 ms to 1143 ms (measured, 214 entries, 932 KB):
import jax
jax.config.update("jax_compilation_cache_dir", "~/.cache/jax")
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0.0) # REQUIRED here
The second line is not optional for this workload: the default threshold is 1.0 s and every one of these compilations is far below it, so with the cache directory alone nothing is ever written. jNO does not set either of these for you — a library should not start writing to a user's disk uninvited.
When AMG is worth it
jno.precond.amg has the best asymptotics on offer: Jacobi's iteration count grows as √n (79 → 166
→ 288 at n = 3k → 12k → 47k on a 2-D Poisson), AMG's is O(1). Build the hierarchy once and reuse
it — an unbuilt spec re-runs pyamg's host-side setup on every solve, and stays off the compiled path
because that setup cannot be traced:
M = jno.precond.amg().build(fem.operator[0]) # or jno.precond.amg().cached()
u = fem.solve(linear=jno.solve.cg(tol=1e-10), precond=M)
Per solve, against cg + jacobi at the same tolerance — the advantage grows with the problem, which
is the √n-vs-O(1) law showing up directly:
| DOFs | cg + jacobi | cg + amg (built) | setup | break-even | |
|---|---|---|---|---|---|
| 3,013 | 6.4 ms | 4.4 ms | 1.5x | 135 ms | 66 solves |
| 18,289 | 13.6 ms | 6.6 ms | 2.1x | 317 ms | 45 solves |
| 46,677 | 31.5 ms | 10.0 ms | 3.2x | 303 ms | 14 solves |
| 95,061 | 80.3 ms | 16.4 ms | 4.9x | 419 ms | 7 solves |
So AMG is for repeated solves against the same operator — a transient run, a parameter sweep, a Newton loop — where the setup amortises. For a single one-shot solve below ~100k DOFs, Jacobi still wins on wall clock: you would pay 419 ms of setup to save 64 ms. This is why AMG is not the default; the right choice depends on how many times you solve, which only you know.
Large systems on a GPU
Past a few hundred thousand DOFs, what stops a GPU solve is rarely the matrix. Three things bite first, and only the third is about size at all.
Assembly runs on the host, automatically. jNO's element loop allocates temporaries far larger
than the matrix it produces, so assembling on the device runs out of memory at roughly a third of the
size the finished operator would happily occupy. jno.fem(...) therefore builds on the CPU whenever
a GPU is the default backend, and the solve moves the operator across:
fem = jno.fem([...]) # assembled on the host, whatever the default backend
u = fem.solve(linear=jno.solve.fgmres(), precond=M) # solved on the GPU
Nothing is given up for this. jax.device_put across backends is traceable and jax.grad
differentiates through it, so topology optimisation, neural coefficients and trainable coordinates
all still work across the split. Measured on a 9,970-DOF Poisson, assembly now costs 8 bytes of
device memory — flat from n=242 to n=3,797 — against 104 MB before. Budget roughly 6 GB of host RAM
per million DOFs; host memory, not the card, becomes the binding constraint.
It is not slower for a one-shot build, either — the host wins on every first call, because it also compiles faster:
| assemble on host | on device | |
|---|---|---|
| Poisson, 43k DOFs | 4.02 s | 4.98 s |
| mixed N1E x Lagrange, 43k DOFs | 9.28 s | 11.00 s |
The device does win on repeated assembly of a heavy form once compilation has amortised (5.60 s on
the host against 3.69 s on the device, mixed N1E at 43k DOFs), which matters in a Newton or
optimisation loop that reassembles many times. There is no jNO argument for this — jax.default_device
is JAX's own way of saying where work goes, and an explicit one wins:
with jax.default_device(jax.devices("gpu")[0]):
fem = jno.fem([...]) # assemble on the device instead
Set the async allocator. XLA's default BFC allocator fragments and will refuse an allocation about a gigabyte below the card's actual free memory — on a GPU shared with a desktop this is the difference between running and not:
A 2.70 GB operator that BFC rejected ran without complaint under cuda_async.
Sizing. A BCOO operator costs 16 bytes per nonzero in float64 (an 8-byte value plus two
int32 indices), and peak device memory during a solve is about 1.3x that, the extra being SpMV
scratch. Measured on an RTX 3070 (7.81 GiB usable, ~3 GB held by a desktop), on the fused real
2n block of a complex mixed N1E x Lagrange system at ~50 nonzeros per row:
| DOFs | nonzeros | operator | peak | SpMV | effective |
|---|---|---|---|---|---|
| 1,215,312 | 60.3 M | 0.90 GB | 1.24 GB | 7.9 ms | 114 GB/s |
| 2,216,184 | 113.3 M | 1.69 GB | 2.22 GB | 15.0 ms | 113 GB/s |
Bandwidth holds flat across sizes, so a Krylov solve costs about iterations x SpMV: budget from the iteration count, not from the DOF count.
float32 is not a lever. The same 60.3 M-nonzero operator takes 8.05 ms in float64 and
9.09 ms in float32 — half the value bytes, and slightly slower. Indices stay int32 either way
and the gather of x dominates, so halving precision buys about a quarter of the memory and none of
the time. Reach for a coarser mesh or a better preconditioner instead.
Preconditioners (jno.precond, for the iterative solvers):
| Kind | jno.precond |
|---|---|
| Diagonal / polynomial | jacobi, chebyshev |
Between jacobi and multigrid |
nystrom (randomized low-rank, SPD); fsai (factored sparse approximate inverse, SPD — applied as two sparse products, numeric setup in JAX; see FSAI) |
| Multigrid | amg (algebraic, pyamg), jaxamg (GPU AmgX), gmg (geometric — a structured-grid V-cycle), ams (H(curl) auxiliary-space Maxwell), hypre (AMS / BoomerAMG via PETSc) |
| Domain decomposition | schwarz (overlapping Schwarz on a METIS partition, two-level by default; needs the [metis] extra — see Overlapping Schwarz) |
| Incomplete factorization | ilu (SuperLU spilu, complex-capable) |
| Block / saddle point | block_diag / triangular (block / Schur), saddle, lsc, pcd |
| Composition | form (weak-form auxiliary operator), inner (any solver as M⁻¹), real_equivalent (a complex system through its real block), cached |
| Precision | float32=True on every constructor but cached — built and applied in single precision while the solve stays in double; see Single precision |
Process-wide solver settings (jno.setup, or the same keys under [jno] in .jno.toml):
lu_stack=k stacks k systems per call when a batched jno.solve.lu() runs on the default device
backend (default 1); matvec_format="auto" | "csr" | "coo" fixes the sparse storage the iterative
solvers apply an operator with (default "auto", timed on the running device per operator size class). See the jno.setup docstring.
jno.solve
jno.solve -- the callables-only solver namespace for fem.solve's slots.
Every factory returns a configured callable (no strings anywhere): a linear solver
(A, b, *, M=None, x0=None) -> x or a nonlinear driver (residual_fn, u0, *,
linear_solve=None) -> u. All shipped solvers are pure JAX -- jit- and vmap-native,
differentiable through lax.custom_linear_solve / lax.custom_root -- and they reuse
existing implementations (jax.scipy.sparse.linalg Krylov, the differentiable sparse-direct
spsolve) rather than re-implementing them. A user-written callable with the same signature
drops into the same slot; if it is pure JAX it inherits every transform automatically.
Usage::
fem.solve(linear=jno.solve.cg(tol=1e-10), precond=jno.precond.jacobi())
fem.solve(linear=jno.solve.lu()) # differentiable sparse-direct
fem.solve(nonlinear=jno.solve.newton(), x0=u_guess) # warm-started Newton
Defaults when a slot is None: Jacobi-preconditioned BiCGStab (steady linear) and :func:newton --
Newton on the assembled tangent with a Jacobi-BiCGStab inner solve, matrix-free Newton-Krylov where the
problem has no assembled tangent (nonlinear).
LinearOperator
Uniform handle over an assembled operator (BCOO, dense array, or bare matvec).
Gives every linear solver one interface regardless of the storage the assembler produced:
.mv(v) (also @), lazy transpose .T, .diag(), .dense(), and .bcoo
(None when not sparse). A matvec-only operator (from_matvec) supports mv and a
transpose via jax.linear_transpose; diag/dense raise -- direct solvers and
diagonal preconditioners need an assembled matrix.
from_matvec
classmethod
from_matvec(mv: Callable, *, t_mv: Optional[Callable] = None, diag_fn: Optional[Callable] = None, dense_fn: Optional[Callable] = None, shape: Optional[tuple] = None, _transposed: bool = False) -> 'LinearOperator'
Wrap a bare matvec. Optional hooks upgrade it: t_mv (transposed matvec — else
derived exactly via jax.linear_transpose), diag_fn/dense_fn (else those
accessors raise), shape (else None).
LinearSolver
LinearSolver(fn: Callable, *, name: str, traits: Optional[dict] = None, direct: bool = False, key: Any = None, settings: Optional[dict] = None)
A configured linear solver: solver(A, b, *, M=None, x0=None) -> x.
M is the preconditioner application v -> M^-1 v. A jno.precond.* spec is
accepted too and materialized against A on the way in, so jno.solve.cg()(A, b,
M=jno.precond.jacobi()) works without spelling out materialize_precond. A bare callable
is always taken as the applier, never as a ctx -> applier factory -- as a precond= slot it
would be the latter, and nothing about a callable distinguishes the two, so the ambiguous case
keeps the meaning the Krylov routines already give it. Specs needing eager preparation
(jno.precond.form, which assembles an auxiliary operator) still need fem.solve(precond=...):
a direct call has no owning FEM to prepare against.
fn receives (op: LinearOperator, b, M, x0). traits documents transform support
(vmap: "native" | "sequential" | "no", jit: True | False) so composition layers (and,
later, the auto policy) can pick honestly instead of silently host-looping. jit=False marks
a solver whose iteration cannot run inside a trace -- Chebyshev measures its spectrum bounds
and branches on the measured values, which a tracer has no answer for.
key is the solver's value identity: the constructor arguments that change what it does.
It exists because the compiled slot path (:func:compose_linear_solve_fn) hands the spec to
jax.jit as a static argument, and jax keys its compilation cache on hash. Hashing by
identity would mean fem.solve(linear=jno.solve.cg()) -- the spec written inline, as the docs
themselves write it -- recompiling on every call: measured 0.4 ms against 83.5 ms on a 513-DOF
Poisson solve, i.e. far worse than never compiling. With a key, two equivalently configured
specs are one cache entry.
A spec with no key falls back to identity, and the composer then declines to compile at all. That default is deliberate: a key that omits a parameter would serve a cached solve configured the other way, which is a wrong answer, while no key merely forgoes a speed-up. So a new solver is slow until its key is written, never silently wrong.
NonlinearSolver
NonlinearSolver(fn: Callable, *, name: str, traits: Optional[dict] = None, direct: Optional[bool] = False, config: Optional[dict] = None)
A configured nonlinear driver: driver(residual_fn, u0, *, linear_solve=None, jacobian=None) -> u.
direct=True marks an assembled-Jacobian, sparse-direct Newton: it factorizes the assembled
tangent (jacobian= — a callable u -> BCOO) each step instead of the matrix-free Krylov inner
solve, so it is robust on indefinite/ill-conditioned systems (Taylor-Hood saddles, stiff drag). It
composes only where the assembler provides that Jacobian (native nonlinear FEM / the transient
stepper), which threads it in via jacobian=.
lu
Differentiable sparse-direct solve (JAX spsolve: cuSolver on GPU, native LU on CPU).
Wraps the existing :func:jno.utils.solver.linear.sparse_lu_solve -- robust on the
indefinite saddle-point systems where Jacobi-preconditioned Krylov stalls, reverse-mode
differentiable in the matrix entries and the right-hand side. Direct: ignores x0 and
rejects a preconditioner. jit yes, and vmap -- so jax.jacrev / jax.jacfwd through
a solve -- on every backend. How a batch is solved depends on the backend:
"host","cudss","pardiso"factor once for a batch against one matrix and solve the whole block of right-hand sides in one call -- the batch to pick for Jacobians and sensitivities. Measured (RTX 3070 box, float64, factorising on every call):jacrevover 32 outputs of a 2-D 20k-DOF solve 108 ms against 58 ms for ONE solve, where"device"took 16.7 s."device"(cuSolver's sparse QR) takes one right-hand side, so it factorises every system: one per call by default, orjno.setup(lu_stack=k)stacked per call.
| PARAMETER | DESCRIPTION |
|---|---|
backend
|
WHERE the factorization happens. All three obey the same
Choosing between the last two: pick by the phase your problem repeats. A Newton loop re-FACTORIZES, so PARDISO wins. A shift-invert eigensolve or a constant-operator transient re-SOLVES against one factorization, and there cuDSS is 11x faster per solve (3.5 ms vs 40 ms at lap3d 50^3) and additionally takes a whole block of right-hand sides at once. There is deliberately no
TYPE:
|
host
|
Deprecated alias for
TYPE:
|
reuse
|
keep the factorization for a repeating operator ( Pass
TYPE:
|
dense
Dense LAPACK solve (jnp.linalg.solve) on the densified operator.
O(N^2) memory / O(N^3) time -- the right answer for small systems and coarse
blocks; vmap-native. Direct: ignores x0, rejects a preconditioner.
cg
Conjugate gradients (jax.scipy.sparse.linalg.cg) -- symmetric positive-definite
systems only (Poisson, elasticity, mass matrices). Cheapest per iteration; takes M and
x0. Implicitly differentiable upstream via lax.custom_linear_solve.
bicgstab
BiCGStab (jax.scipy.sparse.linalg.bicgstab) -- general non-symmetric systems.
With precond=jno.precond.jacobi() this reproduces the historic fem.solve()
steady-linear default exactly.
gmres
gmres(*, tol: float = 1e-08, atol: float = 0.0, maxiter: Optional[int] = None, restart: int = 30) -> LinearSolver
Restarted GMRES (jax.scipy.sparse.linalg.gmres) -- non-symmetric systems where
BiCGStab's erratic convergence hurts; memory grows with restart. For an iterative
(e.g. multigrid-with-tolerance) preconditioner a flexible variant (FGMRES) is required --
planned; see plans/fem-solver-api.md.
fgmres
Flexible restarted GMRES (Saad 1993, Alg. 2.2; see
:func:jno.utils.solver.krylov.fgmres) — the outer solver to use when the preconditioner is
itself iterative (an inner Krylov sweep, a multigrid cycle with a tolerance, a block/Schur
recipe with inexact inner solves), which plain GMRES's fixed-M assumption forbids.
Memory: two (restart, n) bases.
minres
MINRES (Paige & Saunders 1975, §5; see :func:jno.utils.solver.krylov.minres) — the
Krylov method for symmetric indefinite systems: Stokes/Biot saddle points, biharmonic
(Argyris/Morley), shifted Helmholtz-like operators. Monotone residual where BiCGStab is
erratic; O(1) memory where GMRES grows with restart. The preconditioner must be
symmetric positive definite even when A is indefinite.
cocg
COCG (van der Vorst & Melissen 1990; see :func:jno.utils.solver.krylov.cocg) — the Krylov
method for complex-symmetric systems, A = A^T and not A = A^H: time-harmonic
eddy-current (A-V), Helmholtz, RCWA. Replacing the Hermitian inner product by the bilinear
x^T y keeps CG's short recurrence, so this costs O(n) memory and one matvec per
iteration where GMRES needs a (restart, n) basis and a restart parameter.
Use it only when the operator really is complex-symmetric — jno.fem gives that after the
V = jw*Vt column scaling of an A-V system, and a Helmholtz form is symmetric as assembled.
On a Hermitian or non-symmetric operator the recurrence is not valid; use minres or
fgmres. Unlike CG on an SPD matrix, COCG can break down (the bilinear form is indefinite);
it then returns its last iterate and the residual guard raises.
chebyshev
chebyshev(*, lmin: Optional[float] = None, lmax: Optional[float] = None, tol: float = 1e-08, maxiter: int = 500, bound_iters: int = 30, lmin_ratio: float = 1.0 / 30.0, safety: float = 1.05) -> LinearSolver
Chebyshev semi-iteration for SPD systems (Golub & Varga 1961; Saad 2003 §12.3,
Alg. 12.1; see :func:jno.utils.solver.krylov.chebyshev_iteration). Inner-product free —
matvecs and AXPYs only, no reductions — so it shines under vmap and on GPU where CG's
dot products serialise. Needs spectrum bounds of M^{-1} A: pass lmin/lmax when
known; otherwise both ends are measured by bound_iters steps of Lanczos (Lanczos 1950,
§II — the extreme Ritz values of the tridiagonal), for the same one-matvec-per-step cost as
the power iteration it replaces. Without the optional :mod:matfree package this falls back
to power iteration for lmax and the lmin = lmin_ratio * lmax guess, which converges
more slowly and, when the true ratio is smaller than assumed, leaves the lowest modes outside
the fitted interval where the polynomial amplifies them.
amg
amg(*, tol: float = 1e-06, maxiter: int = 500, krylov: Optional[str] = 'PBICGSTAB', config: Optional[dict] = None) -> LinearSolver
GPU algebraic-multigrid solve via jaxamg (NVIDIA AmgX wrapped as a JAX primitive).
A self-contained solver: it runs an AMG-preconditioned Krylov iteration (or pure AMG) entirely on
the GPU/device — the on-device counterpart to the host-side pyamg used by jno.precond.amg.
Ideal for large H¹-elliptic systems (Poisson, diffusion, elasticity) and, via
jno.precond.inner(jno.solve.amg(...)), as the smoother inside an outer flexible-Krylov solve
(e.g. the auxiliary nodal solves of a future H(curl) AMS preconditioner). Plain AMG will not
converge on a raw curl-curl (H(curl)) system — that needs AMS on top.
config (a full AmgX-format dict) overrides the convenience args entirely; otherwise a config is
built from tol/maxiter/krylov (krylov=None → pure AMG, else an AMG-preconditioned
krylov Krylov, e.g. "PBICGSTAB"/"GMRES"/"PCG"). Needs an assembled operator
(hands the sparse matrix to jaxamg); errors on a matrix-free operator. Direct-style: takes no outer
precond= (it owns its AMG preconditioner) and ignores x0.
Optional dependency — jaxamg is imported lazily; see :func:_require_jaxamg for the requirements.
Reference: Liu, Fan & Wang, JAX-AMG: A GPU-Accelerated Differentiable Sparse Linear Solver Library for JAX, arXiv:2606.09001 (2026); wraps NVIDIA AmgX (Naumov et al., 2015).
newton
newton(*, damping: float = 1.0, rtol: float = 1e-08, atol: float = 1e-08, max_steps: int = 100, inner_tol: float = 1e-10, inner_maxit: int = 2000, line_search: bool = False, ls_max: int = 25, ls_c: float = 0.0001, direct: bool | None = None, reuse: bool = False) -> NonlinearSolver
Newton root-find, as a configurable slot. Three inner-solve modes:
- default (
direct=None) -- the ASSEMBLED tangent, solved iteratively. Wherever the assembler provides the tangent (native nonlinear FEM, the transient stepper), each step assemblesJ(u)and solves it with thelinear=/precond=slots -- by default Jacobi-BiCGStab -- so an inner iteration is one SpMV instead of two Jacobian-vector products through the element loop, and a preconditioner that needs an assembled matrix (jacobi, a builtamg) composes. Measured on a 3-D P1-div((1+u^2) grad u)problem (RTX 3070): 2.3-3.4x faster than the matrix-free Newton at 10k-87k DOF, same root. A problem with no assembled tangent (a residual-only operator) falls back to the matrix-free mode automatically. direct=False(matrix-free) --J @ vfrom a JVP, inner matrix-free solve (default BiCGStab, or thelinear=slot). No tangent is ever stored, so it is the choice when the assembled tangent would not fit in memory. The previous default.direct=True(sparse-direct) -- factorize the ASSEMBLED tangent each step with a sparse LU (or thelinear=slot's direct solver). Robust on indefinite / ill-conditioned systems -- a Taylor-Hood velocity/pressure saddle, a stiff Carman-Kozeny phase-change drag -- where an iterative inner solve has no saddle-point preconditioner and stalls.
All three are differentiable (implicit differentiation via lax.custom_root).
damping < 1 relaxes each update; line_search=True adds residual-norm Armijo backtracking (up
to ls_max halvings, constant ls_c) so a stiff problem converges without hand-tuning.
reuse=True is lagged-Jacobian Newton: step against the last ASSEMBLED tangent until its
contraction drops below 1/2, then refresh -- fewer assemblies (and, with direct=True,
factorizations) for more, cheaper steps. A step on a stale tangent that does not reduce the residual is
rejected, so it cannot diverge where the fresh Newton would not. It pays wherever assembling the tangent
costs more than a solve -- a long element integrand (a stabilised flow), or a backend that keeps its
factorization (lu(backend="cudss" | "pardiso"); backend="device" refactorizes anyway). Refused
with direct=False, which never assembles a tangent. fem.stats["nonlinear"]["factorizations"]
reports the count. See :func:jno.utils.solver.newton_krylov.newton_direct for the rule and its source.
A transient march's own default Newton is newton(reuse=True), and an eager march carries the
tangent from step to step as well: the tangent of a march changes little between steps, so it is
assembled a handful of times per march rather than at every Newton iteration (measured on a stabilised
3-D flow: once for 8 steps, 0.80 -> 0.30 s per step). An explicit nonlinear=newton() keeps a fresh
tangent per iteration.
picard
picard(*, damping: float = 1.0, rtol: float = 1e-08, atol: float = 1e-08, max_steps: int = 200, inner_tol: float = 1e-10, inner_maxit: int = 2000, line_search: bool = False, ls_max: int = 25, ls_c: float = 0.0001, anderson: int = 0) -> NonlinearSolver
Damped Picard (lagged-coefficient / fixed-point) iteration — pair with :func:jno.lag.
Freeze the troublesome solution-dependent coefficients in the weak form with
jno.lag(...); the linearization of the residual is then the Picard operator (the
lagged system re-solved at each iterate), and this driver iterates it with optional damping.
The classic trade: more outer iterations than Newton's quadratic convergence, but each
linearized system keeps the structure (symmetry, definiteness) that block preconditioners
and multigrid need — e.g. a non-Newtonian Stokes flow whose full-Newton velocity block is
strongly nonsymmetric while its Picard block is a plain symmetric Stokes operator.
Without any jno.lag marker in the residual this is exactly damped Newton. The default
max_steps is higher than Newton's — linear (not quadratic) convergence. line_search=True
adds residual-norm Armijo backtracking (up to ls_max halvings, sufficient-decrease constant
ls_c): essential when the lagged operator's step overshoots from a stiff initial state (a
rigid-plastic cold start whose effective viscosity spans orders of magnitude), where fixed
damping alone either diverges or crawls. See the jno.lag docstring for the inverse-problem
(Picard-adjoint) caveat.
anderson=m (m >= 1; 0 = off) applies Anderson acceleration over the last m iterates
(Walker & Ni, SIAM J. Numer. Anal. 49(4), 2011): each step extrapolates from the history of Picard
updates, which often turns linear convergence into much faster convergence for the same one linear
solve per step. Safeguarded: an extrapolated point that does not decrease the residual retreats toward
the plain Picard iterate, so it is never worse than anderson=0. Costs 2 m vectors of memory and
one thin QR of an n x m matrix per step; m of 3-10 is usual. Gradients are unchanged (implicit
differentiation at the root).
continuation
Parameter continuation for fem.solve(continuation=...): march runtime parameters across
a value sequence, warm-starting each solve from the previous one::
cap = jno.np.parameter((1,), name="cap")
fem = jno.fem([... cap ...]) # built ONCE
u = fem.solve(continuation=jno.solve.continuation(cap=np.geomspace(G / 2000, G, 8)))
One driver, three names for one mechanism: a frequency or material sweep in EM, load stepping in mechanics, homotopy in numerics -- reaching a parameter value the cold solve cannot. Sequences given together are marched zipped, not as a grid, so two coefficients can ramp in step.
keep="last" returns the final solution (homotopy); keep="all" returns the whole family,
(n_values, n_dofs) -- a sweep.
Why this instead of a Python loop that rebuilds the form. The form is traced and compiled once and the parameter arrives as a runtime argument, so an 8-step ramp is 8 solves, not 8 rebuilds and 8 XLA compilations. That is the difference between a continuation being a tool and being the dominant cost of a run.
Fixed (unswept) parameter values are passed as ordinary keywords to fem.solve alongside this
spec; naming one in both places raises rather than silently picking a winner.
Scope, refused by name elsewhere: steady problems (linear or nonlinear), real or
fused-complex. A transient sweep is a plain loop over fem.solve() -- there is no warm start to
carry between independent trajectories, so this driver would add nothing.
staggered
staggered(fields, *, rtol: float = 1e-08, atol: float = 1e-08, max_sweeps: int = 200, inner_steps: int = 20, inner_tol: float = 1e-10, inner_maxit: int = 2000, line_search='backtrack', damping: float = 1.0, ls_max: int = 25, ls_c: float = 0.0001, direct: bool = False, over_relax: float = 1.0, anderson: int = 0) -> NonlinearSolver
Alternate minimization — solve a coupled system one field at a time, sweeping until the full
residual converges. fields is the trial symbols in the order to sweep them::
fem.solve(nonlinear=jno.solve.staggered([u, dm]))
Reach for it when the coupled energy is non-convex in the fields jointly but convex in each
separately — the case where a monolithic Newton has no descent guarantee and diverges outright.
Variational phase-field fracture is the canonical one: the (1-d)^2 |grad u|^2 coupling is quartic
in the pair, while u alone solves a linear elasticity problem and d alone a linear elliptic
one. Alternate minimization turns that into a sequence of convex solves, each decreasing the energy.
Fixed-stress Biot poroelasticity and thermo-mechanical staggering have the same shape.
Algorithm: Bourdin, Francfort & Marigo, Numerical experiments in revisited brittle fracture, J. Mech. Phys. Solids 48 (2000), §3 — as the staggered operator split with a history field, Miehe, Welschinger & Hofacker, IJNME 83 (2010).
The trade is the convergence rate, and it is not small. Alternate minimization converges
linearly where Newton is quadratic, so it can need hundreds of sweeps near a propagating crack —
hence max_sweeps=200. It buys robustness, not speed; on a problem where Newton converges,
Newton is the better choice (Farrell & Maurini, CMAME 312, 2017, compare the two directly).
Sweeping is Gauss-Seidel: each sub-solve sees the updates made earlier in the same sweep, so the
ORDER of fields matters. Every block must be listed — an omitted field's equations would never be
solved, which is rejected rather than silently skipped.
Differentiable in the ordinary way: at convergence the full residual is zero, so the sweep is just a
way of finding that root, and lax.custom_root supplies the gradient from the full Jacobian.
direct=True factorizes each field's assembled diagonal block instead of solving it
matrix-free, and pairs with a linear= slot::
fem.solve(nonlinear=jno.solve.staggered([u, dm], direct=True), linear=jno.solve.lu(backend="pardiso"))
Reach for it when a sub-block is ill-conditioned — near-incompressible elasticity (ν → 0.5) is the
common one. The matrix-free default cannot help there: a precond= spec materializes against an
assembled operator, and a sub-solve is a restriction closure with none, so the block is solved by
unpreconditioned BiCGStab. The trade is that the full tangent is assembled to use one block of
it; a sparsity-caching backend (pardiso/cudss) then pays only the numeric re-factorization
per sweep. On a well-conditioned problem the matrix-free default is cheaper — this is not a
free upgrade, and it is not the default.
Scope: composes through fem.solve(nonlinear=...) on a multifield problem, which is where the
block layout comes from; it has no meaning on a single field and says so. On a system reduced by
periodic ties, a slip elimination or hanging nodes the driver is handed the reduced iterate and sweeps
each field's REDUCED DOFs (the layout is read off the iterate's size).
Groups. An entry of fields may be a LIST of trial symbols, which are then solved together
inside one sweep rather than alternated against each other::
fem.solve(nonlinear=jno.solve.staggered([[v, p], [T]], direct=True))
That is not a convenience. A velocity/pressure pair cannot be swept apart: the pressure block is
the constraint block (no diagonal — the one :func:jno.precond.saddle finds structurally), so
solving p with v frozen is not a well-posed sub-problem. Any flow staggered against a solid
or a temperature therefore has to group its Stokes pair. A bare symbol among lists is its own group,
so [[v, p], T] is legal and [u, dm] keeps meaning exactly what it did — one field per sweep.
A group holding a constraint field is indefinite, so it wants direct=True: the matrix-free
default puts unpreconditioned BiCGStab on a saddle block, for the reason given above — a sub-solve
is a restriction closure and a precond= spec has no operator to materialize against. This is
documented rather than enforced; fem.solve's own saddle warning already fires on that shape.
anderson=m (m >= 1; 0 = off) applies Anderson acceleration to the sweep, which is the
fixed-point map (Walker & Ni, SIAM J. Numer. Anal. 49(4), 2011) -- the standard remedy for the slow,
linear convergence of alternate minimization. Safeguarded: an extrapolated point that does not
decrease the full residual retreats toward the plain sweep's answer. It composes with over_relax
(the over-relaxed sweep is then the map being accelerated). Memory 2 m vectors.
eigs
eigs(*, k: int = 6, which: str = 'smallest', sigma=None, linear=None, precond=None, tol=None, maxiter=None, X0=None)
Generalized symmetric eigensolver K x = λ M x (K symmetric, M SPD). Returns a callable
(K, M=None) -> (λ, X): the k eigenvalues at the requested end (which='smallest' /
'largest') and their M-orthonormal eigenvectors (Xᵀ M X = I). M=None is the standard
problem K x = λ x.
Use it for modal analysis (vibration), buckling, EM cavity/waveguide resonances and photonic band
structure — everything that is Kx=λMx rather than Ax=b. Build K/M as source-less
jno.fem bilinear forms (or via :meth:FEM.eigs).
Three paths, selected by the arguments. With no iterative argument the pencil is reduced
densely (Cholesky M=LLᵀ → jnp.linalg.eigh on L⁻¹KL⁻ᵀ) — exact, and the right answer
when you want the whole low spectrum of a small problem, but O(N²) memory because it
materializes the operator. Passing precond= switches to preconditioned LOBPCG
(:func:jno.utils.solver.eigen.lobpcg_geneigh, Knyazev 2001), which only ever applies K/M as
matvecs and so runs at a scale the dense reduction cannot. Passing sigma= targets the k
eigenvalues nearest the shift — interior modes (cavity resonances, band structure away from the
band edge), which no extremal-end iteration can reach — via the spectral transformation
θ = 1/(λ−σ) (:func:jno.utils.solver.eigen.shift_invert_geneigh, Ericsson & Ruhe 1980 §2); the
inner solves against K − σM default to a once-factorized host sparse LU, and linear= picks a
different inner solver (e.g. jno.solve.amg() when a factorization is too big)::
lam, X = K.eigs(mass=mass, k=6) # dense
lam, X = K.eigs(mass=mass, k=6, precond=jno.precond.amg()) # LOBPCG, matrix-free
lam, X = K.eigs(mass=mass, k=4, sigma=60.0) # the 4 modes nearest λ = 60
sigma= replaces which= (the target is "nearest σ") and needs no precond= — the
transformation is its own preconditioner. The Rayleigh-Ritz runs in the M-inner product, so the
consistent (non-lumped) mass matrix of an ordinary FEM form is handled directly. tol/maxiter
tune the iterative paths (defaults 1e-6 / 200) and are rejected on the dense path, so a
tolerance can never be silently ignored. If the budget is exhausted before tol the LOBPCG
result is NaN-poisoned rather than silently under-converged. The shift-invert path raises
ShiftInvertNotConverged (from :mod:jno.utils.solver.eigen, a RuntimeError) instead -- for
an exhausted budget, and for a shift on an eigenvalue whose inner factorization degenerates -- and
NaN-poisons only under jit/vmap, where nothing can be raised. jNO never fails silently.
Differentiable on both paths: ∂λ/∂θ for simple eigenvalues (degenerate/crossing
eigenvalues make the derivative ill-defined — use the trace of the cluster). The dense path
differentiates through eigh; LOBPCG freezes the converged eigenvector and differentiates the
Rayleigh quotient, which is the same derivative exactly (∂R/∂x = 0 at an eigenvector) — but its
eigenvectors carry no gradient, where the dense path's do.
logdet
Differentiable, matrix-free log det A (symmetric positive-definite) — stochastic Lanczos
quadrature via the optional matfree package. Scales where a direct factorisation cannot; the
key use is Bayesian log-evidence / marginal likelihood of a FEM precision operator. Returns an
unbiased estimate (variance ↓ with samples, bias ↓ with order). See
:func:jno.utils.solver.matfun.logdet.
trace
Differentiable, matrix-free tr A (Hutchinson) or tr f(A) (fun=, Lanczos quadrature) —
e.g. fun=lambda z: 1/z for tr(A⁻¹) (uncertainty / effective degrees of freedom). Optional
matfree. See :func:jno.utils.solver.matfun.trace.
applyfun
Matrix-free f(A)·v — e.g. one exact exponential-integrator step exp(-dt·A)·v with
fun=lambda z: jnp.exp(-dt*z). order bounds the Krylov dimension rather than fixing it:
the iteration stops at the order that has actually converged, so raising it can only help (unlike
the stochastic estimators above, where order is a real request). symmetric=True (default, Lanczos) assumes A = Aᵀ;
symmetric=False (Arnoldi + an eigendecomposition of the Hessenberg with an analytic Daleckii–Krein
derivative) handles a non-symmetric A (advection–diffusion). Both are differentiable and
GPU-capable — the non-symmetric path for any holomorphic fun on a diagonalizable A.
Optional matfree. See :func:jno.utils.solver.matfun.applyfun.
diagonal
Differentiable, matrix-free estimate of the diagonal of A (Hutchinson) or f(A) (fun=,
A symmetric) — the per-DOF field counterpart of :func:trace. The key use is
fun=lambda z: 1/z → diag(A⁻¹), the pointwise posterior variance / uncertainty map of a FEM
precision, plottable on the mesh. Stochastic (variance ↓ samples, bias ↓ order); optional
matfree. See :func:jno.utils.solver.matfun.diagonal.
svd
Differentiable, matrix-free partial SVD — the k largest singular triplets of a possibly
rectangular operator (Golub–Kahan bidiagonalization, 1965). The non-symmetric counterpart to
:func:eigs: use it for POD / reduced-order bases from a snapshot matrix, and for the
ill-posedness of an inverse problem — the singular spectrum of the parameter-to-observable map
says which modes are recoverable at all. A is touched only through its matvec, so it may be the
JVP of a differentiable FEM solve rather than an assembled matrix. depth (default 2k+10)
must exceed k — see :func:jno.utils.solver.matfun.svd for why, and for the clustered-spectrum
caveat. Returns (U, s, Vt). Optional matfree.
lstsq
lstsq(A, b, *, damp: float = 0.0, atol: float = 1e-06, btol: float = 1e-06, maxiter: int = 100000, x0=None)
Differentiable, matrix-free least-squares min_x ‖A x − b‖² for a rectangular A (LSMR) —
the gap left by the square Ax=b solvers. damp adds Tikhonov + damp²‖x‖² (ill-posed /
rank-deficient inverse problems); x0 an initial guess. Real operators; optional matfree.
See :func:jno.utils.solver.matfun.lstsq.
checkpoint
checkpoint(path: str, *, every: int = 500, keep: str = 'last', resume: bool = True) -> CheckpointSpec
Checkpoint a march to disk: fem.solve(checkpoint=...).
A march holds every frame in memory and returns them only when solve() returns, so a run
that dies -- OOM, a kill, a power cut -- yields nothing, however far it got. This writes
frames to path as they are produced and records what the march needs to restart, so a dead
run costs the last partial chunk instead of everything. Three marches write themselves down:
- a transient (
u.t) march, with a built-in scheme -- the default theta, ortime=theta/bdf2/sdirk/rosenbrock/exponential; - a load-path march (
.i(k)withdomain(tau=...)), on its fixed grid or an explicittau=<schedule>; - a moving mesh (a geometry term) with
adapt=jno.solve.remesh(...). Its restart is the tuple a topology rebuild already re-enters the march with,{"start", "old": (points, cells, state, layout), "budget", "carry"}.
The fixed-mesh marches (transient, load path) are compiled loops, and checkpointing runs them as a
host loop of compiled chunks of every steps instead. The arithmetic is the same -- the transient
trajectory is bit-identical on CPU and within 1 ulp on GPU, the load path agrees to round-off (2.8e-17) -- but a chain of
scans is not one scan, so a checkpointed march cannot be differentiated: under jit / grad /
jno.core it raises. Checkpoint the forward run (fem.solve(<param>=value, checkpoint=...)) and
differentiate without it. (Unrelated to jax.checkpoint, which is gradient rematerialisation.)
Every other solve refuses checkpoint= rather than accept it and write nothing: a steady solve, an
adaptive time step, a sharded or solve_fn= integrator, adaptive / arc-length / contact load
stepping, and a fixed-mesh adapt=.
every -- steps between writes. A moving mesh also writes at every rebuild, which is where the
field layout changes and therefore where a restart has to begin anyway.
keep -- "last" (default) drops written frames from memory: a moving mesh's trajectory
loads them from disk on demand, and a fixed-mesh march returns its frames as a numpy.memmap of
<path>/frames.npy. "all" keeps everything resident, and checkpoints purely for
crash-resilience.
resume -- when path holds an unfinished run, continue it instead of starting over. A finished
run (complete in its manifest) is never resumed; delete the directory to redo it. A fixed-mesh
store records what identifies its march (grid, sizes, scheme, initial state, parameters, operators),
and resuming a DIFFERENT march into it raises instead of splicing the two together.
Example::
traj = fem.solve(
nonlinear=jno.solve.newton(direct=True),
adapt=jno.solve.remesh(alpha=1.2, every=1),
checkpoint=jno.solve.checkpoint("runs/ball", every=500),
)
u = fem_heat.solve(checkpoint=jno.solve.checkpoint("runs/heat", every=200)).fn() # transient
| RETURNS | DESCRIPTION |
|---|---|
CheckpointSpec
|
pass as
TYPE:
|
remesh
remesh(*, criterion: Any = None, anisotropic: bool = False, max_dofs: int | None = None, every: int = 5, alpha: float | None = None, metric_field: int = 0, hmin: float | None = None, hmax: float | None = None, theta: float = 0.5, refine_factor: float = 2.0, max_iters: int = 8, tol: float | None = None, eps: float | None = None) -> AdaptSpec
h-adaptivity for fem.solve(adapt=...): change the mesh to follow the solution.
criterion= refines on a traced expression instead of the recovery error estimator --
which is how production AMR codes actually mark, on a physical quantity rather than an error
estimate. It carries no test function; a criterion is a field, not an equation::
ui = u.bind(x=xi, y=yi)
remesh(criterion=jno.np.sqrt(ui.x**2 + ui.y**2)) # gradient / shock detector
remesh(criterion=phi * (1.0 - phi)) # phase-field interface
remesh(criterion=jno.np.abs(uy.x - ux.y)) # 2-D vorticity, on a vector field
remesh(criterion=d._by_region({"weld": 1.0, "plate": 0.0})) # refine one material
The expression is assembled against this problem's own test function, normalised by the lumped
mass to a nodal field, and integrated per cell; theta, refine_factor and the remesh
mechanism are unchanged, so it composes with everything else here. On a multifield problem
metric_field selects which field's space it is assembled in. A term that already carries the
test function is accepted too.
Second derivatives in a criterion (a Löhner-style |D2u|/|Du| detector) need order >= 2:
a P1 Hessian is identically zero, so at order 1 such a criterion evaluates to nothing.
A criterion may instead be a per-CELL geometry quantity -- d.cell_aspect(),
d.cell_volume(), or anything built from them. It carries no trial or test function, which is
how it is told apart, and it is evaluated rather than assembled. It must be one value per cell;
reduce a multi-component one yourself (jno.np.min(d.cell_angles(), axis=1)).
A criterion may be a CONDITION rather than a ranking -- jno.le(expr, bound) /
jno.ge(...). Then its signed margin marks every cell that breaks it (not a theta fraction,
and theta is refused), and the march stops when none does. This is what lets a mesh condition
be its own trigger, with no cadence or threshold argument::
remesh(criterion=lambda d: jno.le(d.cell_aspect(), 2.0)) # keep every element decent
Measured on a deliberately stretched mesh: worst aspect 2.87 -> 1.57 in one round, 0 marked on the
next. Set a bound the mesher can reach -- an unstructured 2-D mesh bottoms out near 1.2-1.5, and a
tighter bound never settles. A bare comparison (q > 2.0) is refused: it says which cells are
bad but not by how much, so marking would take a fraction of them and leave the rest.
Pass a callable (criterion=lambda d: ...) for a geometry criterion. A geometry node
captures the cell table when it is constructed, so a single node keeps answering for the mesh it
was built on; once refinement changes the topology it is refused by name rather than read as a
shape mistake.
On a steady problem this is the refine loop — solve, estimate (Zienkiewicz–Zhu), mark
(Dörfler theta), refine by refine_factor, repeat up to max_iters — growing the mesh
toward convergence. On a transient problem it remeshes every every steps and carries the
state across (basis-aware transfer), so the mesh tracks a moving feature. It holds a constant
budget -- max_dofs, else the initial DOF count -- on both paths: the isotropic one
refines the marked cells by refine_factor relative to the rest and then scales the whole size
field to the budget, so the wake coarsens instead of the mesh ratcheting up::
fem.solve(adapt=jno.solve.remesh(anisotropic=True, max_dofs=6000, every=4))
fem.solve(adapt=jno.solve.remesh(criterion=1.0 - phi * phi, every=4)) # follow an interface
On a march a criterion= is evaluated on the live state, at the remesh time, at every remesh.
A condition criterion is a trigger there: the mesh is rebuilt only when some cell breaks it, so
remesh(criterion=lambda d: jno.le(d.cell_aspect(), 3.0), every=1) remeshes exactly when the mesh
degrades, and never while the condition holds.
anisotropic=True refines on a Hessian metric (stretched elements aligned to the curvature of the
solution -- or of the criterion, when one is given) instead of isotropic ZZ marking — far fewer
DOFs for a layer or a front, and the right choice for an interface. hmin/hmax bound the edge
sizes; metric_field picks which coupled field drives the metric.
max_dofs counts DOFs -- every unknown of the assembled system, fem.dofs: a vector field
counts each component, a Taylor-Hood P2/P1 pair about nine per vertex in 2-D, a complex field its
real and imaginary halves. The mesher is steered by a vertex count and honours it only approximately,
so every remesh is counted before it is applied and redone with a corrected target if it misses: the
DOF count stays within 20 % of max_dofs (a march holds it there; a steady round is capped at
it), and a remesh that cannot be brought under 1.2 * max_dofs -- an hmax too small to coarsen
that far, say -- raises instead of proceeding over budget.
alpha= is the moving-mesh form, and it is a different operation: instead of meshing the
geometry afresh it re-triangulates the NODES the motion has carried and keeps the triangles whose
circumradius is below alpha × the starting mean edge length — the alpha shape, the remeshing step
of the Particle Finite Element Method (Idelsohn, Oñate & Del Pin, IJNME 61 (2004) 964–989;
alpha shapes: Edelsbrunner & Mücke, ACM TOG 13 (1994) 43–72). Every node stays where it is, so a
P1 state carries across by identity — and because only the elements are re-decided, the topology
may change: two bodies whose gap closes below about 2·alpha·h become one mesh::
fem.solve(adapt=jno.solve.remesh(alpha=1.2, every=1)) # droplets in a void: they merge on contact
Merging is therefore a mesh-length contact model, not film drainage. Requires a geometry term (nothing else moves the nodes), 2-D, and P1 fields.
Steady-only: max_iters, tol, eps (a relative-change plateau detector, not a certified
bound). Transient-only: every, metric_field. Moving-mesh-only: alpha.
| PARAMETER | DESCRIPTION |
|---|---|
anisotropic
|
Hessian-metric refinement instead of isotropic ZZ + Dörfler marking.
TYPE:
|
max_dofs
|
DOF budget (
TYPE:
|
every
|
Transient only — remesh every
TYPE:
|
alpha
|
Moving meshes only — re-triangulate the moved NODES and keep the triangles smaller than
TYPE:
|
metric_field
|
Transient multifield only — index of the field driving the metric.
TYPE:
|
hmin
|
Smallest allowed edge length (default: mean edge / 50).
TYPE:
|
hmax
|
Largest allowed edge length (default: 2 × mean edge).
TYPE:
|
theta
|
Dörfler bulk-marking fraction (0..1).
TYPE:
|
refine_factor
|
Local edge-size reduction applied to marked cells each round.
TYPE:
|
max_iters
|
Steady only — maximum refine-solve rounds.
TYPE:
|
tol
|
Steady only — stop once the global error estimate falls below this.
TYPE:
|
eps
|
Steady only — stop once the round's figure of merit stops moving by more than this (two consecutive rounds required).
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
AdaptSpec
|
The adaptation spec to pass as
TYPE:
|
See :func:relocate for the fixed-connectivity (r-adaptive) alternative: it keeps the topology, so
there is no mesh schedule to freeze and no cross-mesh transfer, and its vertex map is differentiable
in the monitor. The two compose — remesh first, then relocate on the result.
refine
refine(*, criterion: Any = None, theta: float = 0.5, max_iters: int = 8, max_dofs: int | None = None, tol: float | None = None, eps: float | None = None, metric_field: int = 0) -> AdaptSpec
h-adaptivity by local refinement for fem.solve(adapt=...): split the marked cells.
Beside :func:remesh rather than a flag on it, because it is a different algorithm. remesh
rebuilds the mesh at a finer size field, so it needs a geometry to rebuild from and returns a mesh
that does not nest inside the old one. refine splits each marked cell into 4 (a quadrilateral)
or 8 (a hexahedron): local, needs no mesher, works on a mesh loaded from a file, and every existing
node survives with its value::
d = jno.shape.rect(0, 0, 1, 1).quad().structured(n=8).domain()
u = fem.solve(adapt=jno.solve.refine(theta=0.4, max_iters=4))
d = jno.shape.box(0, 0, 0, 1, 1, 1).structured(n=4).quad().domain() # hexes
u = fem.solve(adapt=jno.solve.refine(criterion=jno.np.abs(ui.x))) # composes with a criterion
For hexahedra this is the only h-adaptivity there is. No general all-hex mesher exists -- gmsh's
Recombine3DAll on a plain box returns tetrahedra and no hexahedra -- so remesh has nothing
to remesh to and refuses by name.
The price is that the mesh stops being conforming: a split cell's edge midpoint is not a vertex of
its unrefined neighbour, so its value is not free. Such hanging nodes are constrained to the
coarse facet they lie on, u = sum_i w_i u_parent_i -- the same relation a periodic tie and a
mortar coupling impose, and carried by the same prolongation. Neighbours are kept within one
refinement level (a 2:1 balance), so no constrained node ever has a constrained parent.
theta (Dörfler marking), criterion, max_iters, max_dofs, tol and eps mean
exactly what they do on :func:remesh, except that a split has no size to steer, so max_dofs only
stops the loop and the last split may overshoot it. There is no refine_factor: a split halves the
cell by construction. There is no anisotropic: the split is isotropic, so there is no direction to
stretch along -- that needs a simplex mesh and remesh(anisotropic=True).
Limitations, measured. Quadrilateral and hexahedral meshes only; a simplex mesh refuses by name
and should use :func:remesh, whose mmg path is local already. A hanging node landing on a tied or
periodic interface is refused rather than composed. Steady problems only -- the transient driver
transfers state across a remesh, and that path does not yet carry a constraint set. Geometry is
exact for affine cells; a warped hexahedron's faces are non-planar, so refining it moves the volume
by an O(h^2) amount (measured: 3.9e-04 for a 0.06 warp on a 0.25 cell).
| RETURNS | DESCRIPTION |
|---|---|
AdaptSpec
|
The adaptation spec to pass as
TYPE:
|
relocate
relocate(*, objective: Any = 'equidistribution', method: str = 'descent', max_iters: int = 8, lr: float = 0.003, quality_floor: float = 0.1, relax: int = 60, relax_step: float = 0.1, every: int = 5, escalate: float | None = None, escalate_growth: float = 1.2) -> AdaptSpec
r-adaptivity for fem.solve(adapt=...): move the mesh vertices, keep the connectivity.
Moves the vertices tagged ``domain.variable(region)[i].trainable()`` so the mesh **equidistributes**
the solution's features, at fixed connectivity and no new DOFs::
xm, ym, _ = domain.variable("core", where=interior, split=True)
xm.trainable(); ym.trainable() # BEFORE jno.fem(...)
u = fem.solve(adapt=jno.solve.relocate())
Requires at least one coordinate tagged ``.trainable()`` before ``jno.fem`` (else it raises).
Tagging is **literal and per-axis**: ``xm.trainable()`` frees only the x column. That is the lever for
boundary vertices — free an edge's *along-edge* axis and its nodes slide within the wall; leave the
normal axis untagged and the domain shape is preserved exactly.
**Two objectives, and the choice is the problem's, not a preference.** ``objective=`` picks *what*
descent minimises:
- ``"equidistribution"`` (default) equidistributes an **arclength monitor** — it targets *resolution*,
and wins where a feature is under-resolved or moving.
- ``"energy"`` descends the **Ritz functional** ``J(u_h) = 1/2 a(u_h, u_h) - l(u_h)`` of the form's own
weak terms. Galerkin orthogonality makes ``J(u_h) - J(u) = 1/2 ||u - u_h||_E^2``, so lowering it lowers
the **energy-norm error** directly. With no load it is the Dirichlet energy (up to the factor 2).
Measured on ``-lap u = 1`` on an L-shape: the error fell to 0.76x at fixed DOFs, where the default
raised it to 1.31x. (It used to descend ``a(u_h, u_h)`` whatever the load; with a source that is
``-2 J``, and it made the same error 6x WORSE -- #114.) It exists only for a **linear, symmetric,
single-field steady** form and is refused by name otherwise (nonlinear, transient, complex,
coupled/saddle, advective).
**Integrate the load accurately.** ``J`` is the DISCRETE functional, with the load at the form's
quadrature, and moving vertices changes that quadrature's error too. On a smooth bump source at the
default degree the descent lowered ``J`` by inflating the quadrated load, and landed ABOVE the
converged energy -- an "improvement" that was a quadrature artifact. ``jno.fem(..., quad_degree=4)``
removed it (error 0.70x, measured; 8 gives the same). A load the quadrature integrates exactly
(constant / polynomial of the right degree) is safe at any degree.
- ``"huang"`` is Huang's equidistribution–alignment functional (see :class:`AdaptSpec`).
**Or a weak-form expression**, when the mesh has a job the three functionals cannot state. They are
mesh-*quality* measures: they see the solution only through a monitor, so they can ask for
resolution but not for a physical condition. An expression is assembled exactly as ``criterion=``
is and summed to a scalar, over a **volume or a boundary** region::
xs, ys, ns = domain.variable("side", normals=True, split=True)
ys.trainable() # the wall may move along y only
us = u.bind(x=xs, y=ys)
fem.solve(adapt=jno.solve.relocate(objective=jno.np.inner(us, ns) ** 2))
That is a **free surface**: the wall is moved until the flow through it vanishes. The facet normals
are rebuilt from the moving vertices, so ``n`` is the current mesh's normal, not the initial one.
The gradient runs through the solve, as it does for the strings — matched to central differences at
7.5e-09 on a Stokes channel whose no-slip bottom couples the flow to the wall's position, where the
through-flow falls 11.4x over 60 rounds (12.5x at 120: this is a descent, not a root-find).
h-adaptivity only when r-adaptivity is not enough. escalate=tol adds a fallback: after each
relocation the P1 interpolation error of the relocated mesh is measured, and if its p90 exceeds
tol (relative to the field's range) the march stops moving nodes and adds some, via the
anisotropic mmg path so the new elements are stretched along the feature::
fem.solve(adapt=jno.solve.relocate(method="monge_ampere", every=20, escalate=0.5))
Relocation gets first refusal because it costs ~0.4 ms against the 8-15 s a node-set change costs,
and escalation then *measures* whether that was enough instead of assuming it. The DOF budget
grows by ``escalate_growth`` each time, capped by ``max_dofs``; the loop is self-limiting, since
more vertices lower the indicator and the gate stops tripping.
The trigger is deliberately **not** a shape-quality floor. On an anisotropically adapted mesh an
isotropic quality measure is anti-correlated with the mesh being good: measured on the 600 W melt
ball at 26 / 74 / 163 nodes, the interpolation error's p90 falls 0.306 -> 0.172 -> 0.078 while the
number of cells rejected by ``4 sqrt(3) A / sum l^2 < 0.2`` RISES 1 -> 5 -> 8. Triggering on the
latter refines a mesh that is already right -- and at the final frame it flagged 5 cells of which
**none** was genuinely bad (stretched *and* across the feature), while 1 truly bad cell went
unflagged.
Two things to know. The objective is a **scalar**, so it needs a scalar test function: on a
velocity/pressure saddle the pressure test is picked automatically. And when the expression reaches
its region only through a **bound view** (``u.bind(x=xr, y=yr)``, which absorbs its coordinates),
the test function cannot be auto-bound — carry it yourself, ``objective=<expr> * v_r[0]``. That
case raises with this instruction rather than a trace-level binding error.
The region's facet quadrature tables are built only when the **form** carries a surface term, so a
surface objective needs the boundary term to be in the ``jno.fem([...])`` list (a traction-free
wall, ``0.0 * v_r[0]``, is enough).
Neither dominates, measured on the two problem types the test suite pins:
========================== ========================= =========================
objective L-shape corner (fixed) Allen–Cahn front (moving)
========================== ========================= =========================
``"energy"`` **55 % error cut** 10.7x WORSE than uniform
``"equidistribution"`` 12 % worse **0.51x uniform**
========================== ========================= =========================
So: a **fixed singularity** wants ``"energy"``; an **under-resolved or moving front** wants the
default. The energy is also not scale-free — a vector field carries the energy of all its components,
so ``lr`` is problem-scaled — where the monitor functionals are.
**Two methods.** ``"descent"`` (default) walks the vertices down the equidistribution defect of an
arclength monitor, evaluated *through the differentiable solve*, with a backtracking ``det J`` line
search — on a stiff problem neither a stock optimiser nor an energy barrier can guarantee validity from
outside the step control. ``"monge_ampere"`` instead solves ``m·det(I + H(φ)) = θ`` for a mesh potential
and takes ``x = ξ + ∇φ`` (McRae, Cotter & Budd, *Optimal-transport-based mesh adaptivity on the plane
and sphere using finite elements*, SIAM J. Sci. Comput. **40**(2) (2018) A1121–A1148, arXiv:1612.08077,
§3.1); the displacement is a gradient, so the *whole* map cannot fold and no line search is needed.
Measured on the Allen–Cahn front the suite uses (``h = 0.06``, ``eps = 0.03``, 377 nodes), error on a
common fine grid so the comparison does not depend on where each mesh puts its nodes:
================== =========== ============ ==================
method rel-L2 vs uniform min element quality
================== =========== ============ ==================
uniform 1.096e-01 1.000 0.834
``"descent"`` 3.951e-02 **0.361** 0.503
``"monge_ampere"`` 8.879e-02 0.811 0.160
================== =========== ============ ==================
So descent stays the default: Monge–Ampère converges in far fewer rounds (3–6 against 30) and reaches a
comparable equidistribution defect, but it degrades element quality badly here and the answer with it.
Lowering ``relax_step`` recovers part of the gap (0.811 → 0.633 at ``relax_step=0.02``).
Works in **2D and 3D**, on a scalar or vector field of **any nodal-Lagrange order**, and across linear,
nonlinear, transient, periodic and complex problems (all but complex-*transient*). It does not compose
with a moving mesh (``coord.d(t) - v``) — that driver owns the march.
Further limits, measured rather than argued:
- **The monitor reads vertex values only**, whatever the element order, so at P2 and above it adapts to
the P1 sub-sampling of the field rather than to everything the field resolves. Higher order still
relocates correctly; it just does not get a sharper monitor for the extra DOFs.
- **Monge–Ampère's non-folding is a property of the whole map.** Holding a subset of vertices truncates
it, and the truncation is what can tangle: on a 21² square with a diagonal front, freezing the whole
boundary reached ``min det J = -1.2e-03`` where the full map stayed positive throughout. Freeing
tangential axes recovers nearly all of it. Either method checks ``det J`` each round and keeps the
last valid mesh, so a bad tagging costs accuracy, not correctness.
- Its relaxation is explicit in ``relax_step``: past the stability limit more iterations make things
*worse* (spread 0.111 → 0.292 going from ``relax=40`` to ``300`` at ``relax_step=0.2``).
- The monitor is arclength-based, which suits an **under-resolved** feature; on an already
well-resolved mesh a curvature monitor wins. Not yet selectable.
- Relocation beats :func:`remesh` when features are few and sharp; loses when they are spread through
the domain (with four separated fronts the crossover moved below one element per feature width) or
when the mesh already over-resolves them. The two **compose** — remesh, then relocate on the result.
Args:
method: ``"descent"`` or ``"monge_ampere"``.
max_iters: Outer relocation rounds.
lr: ``"descent"`` only — base step for the RMS-normalised descent.
quality_floor: ``"descent"`` only — a step is halved until no element's ``|det J|`` falls below this
fraction of the initial worst element.
relax: ``"monge_ampere"`` only — relaxation iterations per round (McRae et al. eq. (3.7)). Each is
one Poisson solve against a matrix factorized once for the whole run, so these are cheap.
relax_step: ``"monge_ampere"`` only — the relaxation pseudo-step ``Δt``.
Returns:
AdaptSpec: The adaptation spec to pass as ``fem.solve(adapt=...)``.
enrich
enrich(*, criterion: Any, theta: float = 0.5, max_iters: int = 8, max_dofs: int | None = None, tol: float | None = None, eps: float | None = None, metric_field: int = 0)
p-adaptivity for fem.solve(adapt=...): raise the polynomial order where it is needed,
leaving the mesh alone.
The fourth adaptivity beside :func:remesh (rebuild finer), :func:refine (split cells) and
:func:relocate (move nodes). This one changes neither the points nor the connectivity: it
switches interpolation covers on at the marked nodes, so the field gains coefficients where
the solution needs them and stays P1 elsewhere. Requires a field declared space="cover"::
u, phi = d.fem_symbols(space="cover")
fem = jno.fem([...])
fem.solve(adapt=jno.solve.enrich(criterion=jno.np.sqrt(ui.x**2 + ui.y**2)))
Why this is the cheap route to variable p: enrichment rides the partition of unity, so an
enriched node next to an unenriched one blends automatically. There are no constraint equations
at an order interface and no edge-mode bookkeeping, which is what a hierarchical p-basis needs.
Compose with :func:refine across successive solves for hp — h where the solution is rough,
p where it is smooth.
A criterion is required, and that is a deliberate refusal to guess. Two built-in estimators were measured against a hand-written one and both lost::
driver L2 x DOFs (lower is better)
Zienkiewicz-Zhu recovery 2.74 and ANTI-CORRELATED with the error: it rose
1.3353e-01 -> 1.3377e-01 over eight rounds while
the true L2 fell 4.692e-03 -> 2.677e-03
hierarchical residual 2.67 honest (it vanishes where enrichment is already
on, by Galerkin orthogonality) but no better at
marking, and it left a sharp spike untouched
criterion=sqrt(ui.x**2+ui.y**2) 1.37
ZZ recovers its gradient from the VERTEX VALUES, and enrichment lives in the cover coefficients, so it cannot see the space it is estimating. Rather than pick a heuristic on the caller's behalf -- gradient of which field, reduced how, for a coupled system? -- this asks for the one thing only the caller knows.
Example -- a gradient-magnitude criterion, the general-purpose choice::
d = jno.shape.rect(0.0, 0.0, 1.0, 1.0, size=0.02).domain()
d.tag("walls", lambda x, y: (x < 1e-9) | (x > 1 - 1e-9) | (y < 1e-9) | (y > 1 - 1e-9))
co, cw = d.variable("interior", split=True), d.variable("walls", split=True)
u, phi = d.fem_symbols(space="cover") # the enriched space p-adaptivity needs
ui, vi = u.bind(x=co[0], y=co[1]), phi.bind(x=co[0], y=co[1])
fem = jno.fem([ui.x * vi.x + ui.y * vi.y - f * vi, u(cw[0], cw[1]) - 0.0])
fem.solve(
solve_fn, # adapt= does not take linear=; pass it here
adapt=jno.solve.enrich(
criterion=jno.np.sqrt(ui.x**2 + ui.y**2),
theta=0.5,
max_iters=6,
),
)
Note ui.xx is NOT available: the assembler takes gradients of a trial function only, so the
curvature that would be the textbook p-indicator cannot be written. A first-derivative criterion
is the practical stand-in.
| PARAMETER | DESCRIPTION |
|---|---|
criterion
|
Required. A traced expression marking where to enrich, exactly as in
:func:
TYPE:
|
theta
|
Dörfler bulk-marking fraction, over NODES rather than cells: the fewest nodes whose
indicator reaches
TYPE:
|
max_iters
|
Maximum enrich-solve rounds.
TYPE:
|
max_dofs
|
Stop once the system reaches this many ACTIVE DOFs. Active, not total: the padded layout gives every node its cover slots and an unenriched node simply has them pinned, so the total never changes and only the free count tracks the enrichment.
TYPE:
|
tol
|
Stop once the global indicator falls below this. Requires an explicit
TYPE:
|
eps
|
Stop when the indicator plateaus for two consecutive rounds. Same requirement, and the same reason with a sharper edge -- against a frozen estimate the relative change is exactly 0.0, so this would declare a plateau on a loop that is still converging.
TYPE:
|
metric_field
|
Which field of a coupled problem drives the marking. Note a coupled form
carrying a cover field is refused by
TYPE:
|
Calling this twice resumes: the mask is state on the domain, so a second run continues from
the first instead of discarding it, and two budgeted calls land where one run to the larger budget
would. A fresh run starts from plain P1 -- an unmasked space="cover" field is already enriched
everywhere, and a loop that begins maximal has nothing left to select.
Scope: simplices only, first-order covers, and the enriched field must be space="cover".
Scalar and VECTOR fields are both supported, in 2-D and 3-D; a vector field gets vec cover
slots per node and every component of an unenriched node is pinned.
fem.adapt_history records n_enriched per round beside the usual n_dofs/estimate.
theta
θ-method time scheme for fem.solve(time=...): θ=1 backward Euler (default),
θ=1/2 Crank–Nicolson / trapezoidal (2nd-order accurate), θ=0 forward Euler. Overrides the
scheme the assembly picks; composes with linear=/precond= (the per-step solve).
Marches the domain's fixed time grid. Call .adaptive(...) on the result to have the step size
chosen from an error estimate instead — jno.solve.theta(0.5).adaptive(rtol=1e-5) — which is
substantially cheaper per digit than the first-order default (see :func:adaptive).
bdf2
BDF2 time scheme for fem.solve(time=...): the second-order backward differentiation
formula (3u^{n+1} - 4u^n + u^{n-1})/(2 dt), and the workhorse for incompressible flow.
Second-order accurate and L-stable, which the theta-method cannot be at once: theta=1
(the default) is L-stable but only first order, and theta=1/2 (Crank-Nicolson) is second
order but merely A-stable -- its amplification factor tends to -1 as dt*lambda -> -inf,
so a stiff mode does not decay, it RINGS. On a saddle-point system (Stokes / Navier-Stokes),
where the pressure has no time derivative at all, that ringing is exactly what you do not want.
Marches the domain's fixed time grid. The first step is plain backward Euler (a multistep method
has no second level to start from); its O(dt^2) local error is what the second-order global
rate needs.
A state-dependent mass c(u)*u_t (variable density, an apparent heat capacity) is marched
in BDF2's non-conservative form c(u^{n+1})(3u^{n+1} - 4u^n + u^{n-1})/(2 dt), second order in
time (measured on a manufactured c(u) = 1 + u^2). It is not the conservative form an enthalpy
mass H(u)_t would want, which needs H itself rather than c = H'.
Refused, loudly, rather than silently mis-integrated: a second-order-in-time (u_tt) block
(assembled at theta=1/2 so an undamped wave is not damped -- and an L-stable scheme would damp
it). .adaptive() is also refused: the controller sizes a one-step method by step doubling.
Reference: Curtiss & Hirschfelder, PNAS 38 (1952) 235; Hairer & Wanner, Solving Ordinary Differential Equations II, 2nd ed., Sec. V.1.
sdirk
SDIRK time scheme for fem.solve(time=...): a singly diagonally implicit Runge-Kutta method,
stiffly accurate and L-stable -- Alexander's 2-stage order-2 (order=2) or 3-stage order-3
(order=3, default) method (R. Alexander, SIAM J. Numer. Anal. 14(6), 1977).
Third order where theta/bdf2 stop at second: for a smooth solution the same accuracy at far
larger steps. Each stage is one implicit solve of the same kind the backward-Euler march takes (a
linear solve, or a Newton solve for a nonlinear block), all with the SAME step operator
M + gamma*dt*A -- so linear=/precond= slots, a sparse-direct factorisation and a constant
preconditioner are built once. Stiffly accurate, so Dirichlet / algebraic (zero-mass) rows are exact
at every step, and L-stable, so it damps stiff transients the way backward Euler does (unlike
Crank-Nicolson). Composes with .adaptive(...) (step doubling with exponent 1/(order+1)).
Costs order implicit solves per step. A state-dependent mass c(u) u_t is integrated in its
non-conservative form (as bdf2 does). Refuses a second-order-in-time (u_tt) block: an L-stable
scheme would damp the undamped wave.
rosenbrock
Rosenbrock (linearly implicit Runge-Kutta) time scheme for fem.solve(time=...): no Newton loop --
each stage is ONE linear solve with the same matrix M + gamma*dt*J (J the tangent at the step's
start), so a nonlinear problem costs a fixed number of linear solves per step.
"ros34pw2"(default) -- Rang & Angermann, BIT Numer. Math. 45 (2005): 4 stages, order 3, stiffly accurate and L-stable, a W-method (order kept with an approximate tangent) and consistent for index-1 DAEs, so Dirichlet and algebraic (zero-mass) rows are handled."ros2"-- 2 stages, order 2, L-stable (Verwer et al., SIAM J. Sci. Comput. 20(4), 1999); not stiffly accurate, so preferros34pw2with algebraic rows.
Composes with linear= / precond= (the stage matrix is what they see; on a linear block with a
constant operator it is built once) and with .adaptive(...). Refuses what a linearly implicit form
cannot carry: a time-dependent or parameter-dependent mass, a state-dependent mass c(u) u_t, and a
second-order-in-time (u_tt) block -- use sdirk() / bdf2() for the first two.
exponential
Matrix-exponential time scheme for fem.solve(time=...) — advances a linear block
M u̇ + A u = f(t) with time-independent M, A by u(t+dt) = exp(-dt·M⁻¹A) u(t) +
forcing, matrix-free. The homogeneous decay is exact in time and unconditionally stable, so it
takes large stiff steps a θ-step cannot. A constant source rides a φ₁ weight (exact); a
time-varying source f(t) is integrated by ETD2 (the exponential trapezoidal rule) — sampled
at both step ends with a φ₂ ramp weight, so it is exact for a source affine in time and
second-order for a general one (Hochbruck & Ostermann, Exponential integrators, Acta Numerica 19
(2010) 209–286, §2.3). order is the Krylov size. mass='lumped' (default) is the row-sum diagonal
— cheapest, discrete maximum principle; mass='consistent' uses the full M (no lumping error) via
a matrix-free M-inner-product Lanczos.
symmetric=True (default, A = Aᵀ) uses Lanczos. symmetric=False handles a non-symmetric
operator (advection–diffusion / transport): it advances by Arnoldi + a differentiable Padé
exponential, with forcing carried exactly through an augmented generator (a ramp row for ETD2) — still
matrix-free, GPU, and reverse-mode differentiable. All paths are differentiable; time-varying
coefficients M(t)/A(t) (a moving/parametric operator) or a nonlinear form → use
:func:theta.
Walls must be homogeneous. Every zero-mass DOF is held at 0, which is a Dirichlet row with
g = 0 and nothing else. A non-zero or time-varying wall value, or an algebraic row that reaches an
unconstrained DOF (a pressure row), raises NotImplementedError rather than silently returning 0 there
(checked eagerly; a traced call relies on the eager one before it). Lift the boundary data
(w = u - g) or use :func:theta, :func:bdf2 or :func:sdirk.
adaptive
adaptive(*, rtol: float = 0.0001, atol: float = 1e-06, max_steps: int = 1000, dt0: float | None = None, limit=None, shrink: float = 0.5, grow: float = 1.5)
Adaptive step-size time scheme for fem.solve(time=...): the step size is chosen per step
from a step-doubling (Richardson) local-error estimate — one full step compared with two
half-steps — so a stiff / sharp / multi-rate transient takes small steps only where it needs them and
large steps elsewhere, instead of the fixed dt from domain(time=(t0,t1,n)).
Built on the block's own implicit θ-step, so it inherits the DAE (Dirichlet) handling and works for a
linear or nonlinear, scalar or vector, plain, periodic, or complex transient. It is a
fixed-length lax.scan of max_steps attempts (a static trip count — the settled tail just
consumes an attempt), which keeps it reverse-mode differentiable (the gradient flows through the
realized step sequence). rtol/atol set the mixed relative/absolute tolerance; max_steps is
the step budget — if it is exhausted before t1 the trajectory is returned as NaN (raise it),
never silently under-resolved. Composes with linear=/precond= (the per-step solve).
Backward-Euler order (first-order in time); pair with a fine tolerance for accuracy.
dt0 is the first step. It defaults to the smallest allowed step (1e-4 of the time span) so
the controller approaches the right step size from below. That default matters: there is no step
rejection (see :func:jno.utils.solver.timeschemes.adaptive_march — a discarded state would make the
per-step solve adjoint run at zero cotangent and return a NaN gradient), so an over-large step is
committed, not retried, and only the next one shrinks. Growing into the step size can never commit an
over-tolerance step — an under-sized step is wasteful, not inaccurate — whereas starting at the output
grid's dt bakes in the error of the first few steps permanently. On a 2-D heat benchmark, taking
dt0 from the output grid left the result 4.2x less accurate than growing from below, for the same
tolerance and ~18% fewer steps. Pass dt0 explicitly only if you know the correct scale; the growth
cap is 5x per step, so the default reaches any scale in a handful of attempts.
This bare form sizes whatever step the assembly picked for the block — backward Euler for a
parabolic block, which is first order. That is the single biggest accuracy lever here, and it is
worth moving: step doubling costs 3 implicit solves per step, so spending them on a first-order base
cannot beat a first-order fixed march by much. Attach the controller to a second-order step instead
and the same tolerance is far cheaper per digit — the step-size exponent follows the base's order
automatically (1/(p+1))::
fem.solve(time=jno.solve.adaptive(rtol=1e-4)) # 5.1e-3 for 162 solves
fem.solve(time=jno.solve.theta(0.5).adaptive(rtol=1e-4)) # 2.2e-4 for 48 solves
(~23x the accuracy on ~3x less work, measured on the 2-D heat benchmark of
tests/test_fem_adaptive_timestep.py — mesh_size=0.15, t=0.05, x64, against the semidiscrete
reference; step doubling costs 3 implicit solves per step). Every scheme exposing a single
step carries .adaptive(...), so this composes with future base methods without new arguments;
:func:exponential raises, since it is already exact in time for the homogeneous decay.
Second order is opt-in rather than the default because θ=1/2 is A-stable but not L-stable: on a
stiff problem with rough or incompatible initial data it rings instead of damping, where backward Euler
is unconditionally smooth. Prefer jno.solve.theta(0.5).adaptive(...) for smooth parabolic and wave
problems; keep the bare form when robustness matters more than order.
NOTE on what adaptivity does and does not buy: on the benchmarks measured here a well-chosen fixed
dt matches or beats it at equal work, because the optimal step size is nearly constant and the error
estimate costs 3x. Reach for adaptive when you cannot pick dt in advance — unknown or
parameter-dependent stiffness, sweeps, inverse problems whose fitted parameter moves the timescale —
not as a speed optimization.
On a pseudo-time LOAD PATH — fem.solve(tau=jno.solve.adaptive(limit=...)) on a
domain(tau=...) history march — the criterion is different, and it has to be. A rate-independent
load path has no local truncation error to estimate: each step is an equilibrium, not an
approximation to a trajectory, so Richardson measures nothing. limit instead bounds how much the
solution may change in one step::
fem.solve(tau=jno.solve.adaptive(limit=0.05)) # every DOF
fem.solve(tau=jno.solve.adaptive(limit=[(dm, 0.05)])) # per field — the usual case
A step is rejected (and the step size cut by shrink) when the solve fails to converge or the
change exceeds limit; a comfortable step grows by grow. That matters beyond cost: with a
fixed grid a step can converge perfectly and still skip an entire propagation event, giving a valid
sequence of equilibria with no resolved event between them — and because the march is path-dependent
(history + irreversibility), that is a different answer, not just a coarser one.
limit is required in the tau= slot and rejected in time= (and vice versa for
rtol/atol) — the two controllers measure different things and silently applying one where the
other was meant would be a plausible wrong answer.
arclength
Arc-length continuation for the load path — a spec for the fem.solve(tau=...) slot.
Load control cannot pass a limit point: past the peak there is no equilibrium at a higher load, so
no amount of step cutting finds one, and :func:adaptive says exactly that when it hits its floor.
Arc-length advances along the equilibrium path instead of along the load, so the path may turn
around — snap-through and snap-back, post-buckling, a softening damage or plasticity branch.
Nothing in the term list changes. The load is still written as a formula in tau; what changes is
that tau becomes an unknown solved alongside u, subject to Crisfield's constraint on the
increment from the last converged point (Crisfield, Computers & Structures 13 (1981) 55-62;
Non-linear FEA of Solids and Structures Vol. 1 §9.3.2)::
Du . Du + psi^2 Dlam^2 = ds^2
Because tau reaches the residual as a plain argument rather than as a DOF, this adds no unknown
to the operator — the border rides a residual wrapper, the way field.bounds(lo, hi) does.
| PARAMETER | DESCRIPTION |
|---|---|
psi
|
weight on the load term.
TYPE:
|
ds
|
the arc length per step.
TYPE:
|
The declared domain(tau=(start, end, n)) is reinterpreted, and this is a real change of
meaning: start is the first (load-controlled) solve, n is the number of output rows, and
end sets the traversal direction and the default ds. It is not a target — the march takes
n steps and stops wherever the path has reached, which is the point, since where the path goes
is the answer. So n buys resolution, not reach: to follow the path further, widen end;
adding steps at a fixed end re-resolves the same total arc more finely.
The load factors actually reached are recorded on fem.tau_schedule — that, against a reaction
read with fem.eval, is the force-displacement curve::
sol = fem.solve(tau=jno.solve.arclength(ds=1e-3))
load = fem.tau_schedule # non-monotone across a snap-back, by design
force = [fem.eval(momentum, sol[k])[grip].sum() for k in range(len(load))]
Scope, stated up front:
- Matrix-free drivers only. The bordered system goes to an unchanged Newton-Krylov, whose
jax.linearizebuilds the bordered JVP for free.jno.solve.newton(direct=True)andjno.solve.staggered(...)are refused by name — the latter because freezing the load factor while sweeping a block IS load control, which has no equilibrium past the fold. - The trajectory is differentiable; the load factors are not.
fem.tau_scheduleis concrete, for plotting and for reading the curve. A runtime-parametric form is refused for the same reason. - Same assembly scope as the rest of the load-path march: real, steady, native Lagrange.
- Adaptive
dsis not implemented — it would reintroduce step rejection, which is what the pilot/freeze/replay split exists to keep out of a differentiable scan.
contact
Let the contact pairing follow the solution — a spec for the fem.solve(contact=...) slot.
u.gap(secondary, main) precomputes, for every secondary quadrature point, which main nodes it
reads and how far it stands off them. That pairing is built once, from the reference
configuration, and is correct only while displacements stay far below the element size. Past that
a point is still tied to the facet it faced before anything moved::
u = fem.solve(contact=jno.solve.contact()) # re-pair from x + u until it settles
Why this is not a refinement detail. Nothing reports a stale pairing. The solve converges
perfectly well; it just converges for a contact configuration that is not the one being solved, and
only a check like a kinematic oracle exposes it — the animation looks right throughout. On a 12:20
involute gear pair, against |T_B/T_A| = z_B/z_A — which holds because the line of action is
common, so each moment arm is that gear's base radius — as the rim mesh went 0.050 -> 0.018
(6856 -> 16418 DOF)::
frozen pairing 0.84% 0.85% 0.85% 0.85%
re-paired 0.64% 0.64% 0.65% 0.64%
Both are flat, and the search buys about 0.2 points: on a rolling contact that barely slides this
is an accuracy refinement rather than a rescue. (The ~0.6% the two share is NOT the pairing — it is
flat in h across a 2.4x DOF range and non-monotone in the penalty, 0.57/0.64/0.73/0.20% over
C_N = 4e4..4e7, so it is neither discretisation nor contact compliance.)
Where the pairing does go stale the error is gross rather than subtle: a flat-bottomed block slid 0.9 across a disk of radius 1 keeps reporting the 0.05 separation it had at its starting position, where the truth is 0.182 — nearly four times larger, and silent.
Each round solves the ordinary system, then re-runs the search at x + u. It stops when the
pairing is unchanged and the solution has stopped moving; exhausting rounds raises and
names how many slots are still oscillating.
| PARAMETER | DESCRIPTION |
|---|---|
capture
|
search radius.
TYPE:
|
rounds
|
maximum solve/re-pair rounds. Reaching it without settling raises, and says which of
the two conditions failed. The default is generous because a round that settles returns
immediately, so the only cost of a high ceiling is paid by a problem that was going to
raise anyway — whereas a ceiling set just at the edge turns a converging solve into an
error. Measured on the gear pair: the whole loop settles in 3 to 4 rounds, the pairing
freezing a round before the displacement reaches
TYPE:
|
tol
|
tolerance on
TYPE:
|
relax
|
damping on the round-to-round update, Converging is not the same as converging to the right branch, and this argument cannot
tell you which you got. A search that oscillates may be oscillating around the answer
or between a correct and an incorrect configuration; damping settles it on whichever
fixed point it is nearest, and reports success either way. Measured on a 12:20 involute
gear pair at two drive angles where the undamped search raised: damped it converged in 17
to 30 rounds and returned torque ratios of 0.0155 and 0.676 against an exact 1.667 -- the
teeth had come out of engagement -- while the FROZEN reference pairing on the same problem
gave 1.05% and 0.62% error.
TYPE:
|
Scope, stated up front:
-
Either tangent works, and the assembled one is faster here. The matrix-free Newton (
newton(direct=False)) re-pairs for free (its tangent isjax.linearizeof the residual).nonlinear=jno.solve.newton(direct=True)assembles the tangent instead and rebuilds the contact block's sparsity pattern per round — sound because the block's SIZE is fixed by the declaration, so only the index values move. Measured on a 12:20 gear pair, 11 682 DOF, 5 rounds, median of three, when the matrix-free Newton was still the default::matrix-free (direct=False) 44.2 s 1848 MB newton(direct=True) 14.7 s 1894 MB <- 3.0x faster, same memory
Peak RSS is comparable at this size rather than a trade: the assembled contact block is small
next to the ~1.8 GB the JAX/XLA runtime already holds, so its cost does not surface. Expect that
to change as the interface grows — the time gap is the robust part (matrix-free spanned
42.7-48.6 s across runs, direct 14.5-15.1 s). Neither is the right default for every problem,
which is why this stays on the nonlinear= slot that already owns the choice.
* The search is host-side and not differentiable in the mesh coordinates. As with the frozen
pairing, the gap is differentiable in the DOF values; d(pairing)/d(x) does not exist.
Each round is its own solve, so a gradient through the loop is not provided either.
* A load path re-pairs per STEP, and is not differentiable. On a form that marches (step
history plus a domain(tau=...) grid) the march owns the loop and runs the search at every
load step — because the pairing that is right at the end of the path is not the one that was
right in the middle of it. That march is a host loop rather than a lax.scan, so it gives up
the load path's reverse-mode differentiability; the scanned march keeps it, at a frozen pairing.
tau=jno.solve.adaptive(...) and tau=jno.solve.arclength(...) are refused by name: both
replay or root-find under a scan, where a host-side search cannot run.
* The pairing is closest-point; there is no smoothing (no NTS-to-mortar blending, no C1 surface),
so a secondary point crossing a main facet edge moves discontinuously. That is what rounds
iterates on and what the raise reports when it oscillates.
jno.precond
jno.precond -- preconditioner specs for fem.solve(precond=...).
A spec is declarative: it says what preconditioner to build, and jno materializes it at solve
time against a :class:jno.utils.solver.solver_api.PrecondContext (the assembled operator; later
per-field blocks and auxiliary weak-form assembly). The materialized applier is just
v -> M^{-1} v and composes with any Krylov solver from jno.solve.
Preconditioners change convergence speed, never the converged solution, so a spec needs no
gradient path -- arbitrary (even non-JAX, via a future callback tier) appliers stay compatible
with differentiable solves. A user spec is any object with materialize(ctx) or a bare
ctx -> (v -> M^{-1} v) callable::
def my_precond(ctx):
inv = 1.0 / ctx.diag()
return lambda v: inv * v
fem.solve(linear=jno.solve.cg(), precond=my_precond)
Composition: :func:block_diag / :func:triangular build block preconditioners over the
per-field DOF blocks (fem.blocks -- or, on a system reduced by periodic ties / slip / hanging nodes,
the per-field blocks of the reduced P^T A P the solve runs on); :func:form assembles auxiliary weak-form operators
("preconditioners as weak forms"); :func:inner turns any jno.solve solver into an
(inexact) M^{-1} application.
PrecondContext
What a preconditioner spec sees at materialization time.
ctx.A is the assembled :class:LinearOperator (matvec-only on the Jacobian-free
nonlinear path), ctx.fem the owning :class:jno.FEM (None outside fem.solve),
ctx.diag() the operator diagonal. For multifield systems ctx.blocks are the
per-field DOF slices of ctx.A (from fem.offsets, or — when a periodic tie, slip
elimination or hanging-node constraint reduced the system to P^T A P — the per-field
slices of that reduced system), ctx.block_slice(field) resolves a trial symbol (or
integer index) to its slice, and ctx.sub(i, j=None) is the (i, j)
sub-operator as a :class:LinearOperator — applied through the full operator's matvec
(embed into block j, extract block i), so it stays sparse/matrix-free; diag and
dense are exact views for i == j direct/diagonal inner solvers.
ctx.assemble(terms, quad_degree=...) assembles an auxiliary weak form with the
ordinary jno.fem machinery and returns its operator — the "preconditioners are weak
forms" primitive (weighted mass matrices, low-order proxies, shifted operators). On a
reduced system the form (assembled on the full finite-element space) is reduced with the
same P as ctx.A, so the two act on the same space.
grid
property
Structured-grid descriptor {shape, spacing, origin} when the operator lives on a regular
grid (jno.shape.rect(...).structured().domain()), else None — needed by geometric multigrid
(:func:jno.precond.gmg). An explicit override if given, else derived from the owning FEM's
domain (ctx.fem.domain.mesh_connectivity["grid"]).
hypre
hypre preconditioners (AMS, BoomerAMG) via PETSc — a reference-grade H(curl) solver.
One duck-typed entry for the whole library rather than a wrapper per preconditioner: kind is
passed straight to PCHYPRE and keyword options become -pc_hypre_<kind>_<key>::
fem.solve(linear=jno.solve.fgmres(),
precond=jno.precond.real_equivalent(
jno.precond.triangular((u, jno.precond.hypre(kind="ams")),
(p, jno.precond.amg()))))
kind="ams" is the auxiliary-space Maxwell solver (Hiptmair & Xu 2007; Kolev & Vassilevski
2009) — the same method as :func:ams, but hypre's production implementation, which additionally
carries the machinery for a semidefinite problem where sigma = 0 over part of the domain
(interior-node marking, a restricted gradient, gradient projection). jNO's own :func:ams has
none of that. Measured on an A–V eddy problem where it matters, hypre reached 8.6e-03 with
interior nodes marked against no progress at all without them.
It is not automatically faster. On the same 89,920-DOF H(curl) block of a production-scale
A–V system, hypre reached 4.33e-09 and jNO's own ams() reached 8.03e-09 — equivalent. Reach
for this as a reference implementation to check jNO against, or for the sigma = 0 features,
not on the assumption that a production library must win.
Requires petsc4py built with hypre, and runs on the host through pure_callback.
That costs a device round-trip per apply and is not traceable, but it does not affect
differentiability: a preconditioner cannot change the solution, so
:func:jax.lax.custom_linear_solve differentiates through the operator and never through the
preconditioner — the same contract the SuperLU, PARDISO and pyamg paths already use.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
ilu
Incomplete LU (SuperLU spilu) — the classical preconditioner for the A–V eddy-current
system, and complex-capable, so it applies to the complex operator directly.
Keyword options pass through to :func:scipy.sparse.linalg.spilu (drop_tol, fill_factor,
permc_spec, …); the defaults here are drop_tol=1e-4, fill_factor=10. equilibrate=False
disables the symmetric diagonal scaling applied first — leave it on unless you know the operator
is already well scaled, because drop_tol is relative and a matrix whose rows span orders of
magnitude (a gauged A–V system runs from 1.0 on pinned rows to 1e11 on conductor rows) otherwise
fails outright with "Factor is exactly singular".
Why this, on a mixed A–V system. Igarashi & Honma (IEEE Trans. Magn. 38(2), 2002, 565–568)
solve exactly this system — edge elements for A, nodal for V — with incomplete-Cholesky CG
and report 16–17 iterations, independent of frequency (100 Hz down to static) and 29–31
independent of mesh size. Their §IV gives the mechanism: the "floating singular values"
(proportional to omega*sigma*mu) that wreck the condition number disappear after incomplete
factorisation in the A–V method, while they survive in the A-only method. That is the documented
reason A–V converges better than A, and it is a property of the incomplete factorisation, not of
multigrid::
fem.solve(linear=jno.solve.fgmres(tol=1e-9, maxiter=400), precond=jno.precond.ilu())
Do not gauge it. The same paper: "the tree–cotree gauging has not been applied because it is known that the gauging makes the matrix condition worse."
Runs on the host and is not traceable; that costs a callback per apply but has no effect on
differentiability — see :func:hypre for the same contract. Memory is the real limit: spilu
fill grows with fill_factor, so on a large 3-D curl-curl expect to trade drop_tol against
RAM before it beats an auxiliary-space method.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
real_equivalent
Precondition a COMPLEX system through its fused real-equivalent block, with a real inner solver — so a real-only preconditioner (an AmgX-style GPU multigrid, say) works on it.
A complex symmetric A = K + iM becomes [[K, -M], [M, K]] over [x_r; x_i]. That block is
skew-dominated (measured ||A - A^T|| / ||A|| = 2.0 when the mass term dominates), which is why
multigrid applied to it directly diverges. The classical fix is not to precondition the block
itself but to use K + M: real, symmetric, and definite whenever K is and M >= 0. Here
that is applied to each of the two diagonal halves, and inner is any ordinary real spec::
fem.solve(linear=jno.solve.fgmres(),
precond=jno.precond.real_equivalent(jno.precond.amg()))
Measured on a complex A-V (N1E x Lagrange) system, GMRES to 1e-8 with an AMG inner solve: 18 iterations where the imaginary part is definite.
It is only as good as the gauge. On an eddy-current problem M = w sigma VANISHES outside
the conductors, which is exactly the assumption (definite imaginary part) most of this literature
makes. Measured on such a system, regularised by a displacement-current mass term of size eps:
eps = 1e-6 -> 40000 its | 1e-3 -> 2946 | 1e-2 -> 248 | 1e-1 -> 54 | 1 -> 16
i.e. the iteration count tracks the gauge, not the vanishing M — with a well-conditioned gauge
it matches the definite case. An eps mass term buys that conditioning by changing the physics;
a tree-cotree gauge (domain._extra_dof_pins) buys it for free, and is the right pairing.
inner sees K + M and never sees a complex number, so anything real composes -- but it must
MATCH THAT OPERATOR'S CHARACTER, which is the practical trap. Where M vanishes, K + M is a
bare curl-curl operator, and smoothed aggregation is the wrong tool for H(curl); that is what AMS
is for. Measured on a cube: with a mass term everywhere (eps-gauge 1) real_equivalent(amg())
converges in ~20 iterations, and at eps 1e-4 -- so K + M is curl-curl in the non-conducting
region -- the same call fails. Conductor thickness and volume fraction are NOT the variable: 20
iterations from 60 % of the volume down to 0.86 %, and 4 elements thick down to 1.
So on an eddy problem with a non-conducting region, pair this with an H(curl) inner (ams()),
not amg(). K + M being real is what keeps a GPU route open -- it wants a real AMS.
References: Day & Heroux, Solving complex-valued linear systems via equivalent real formulations, SIAM J. Sci. Comput. 23(2):480-498, 2001 (which equivalent real form to pick); Axelsson & Kucherov, Real valued iterative methods for solving complex symmetric linear systems, Numer. Linear Algebra Appl. 7:197-218, 2000; Benzi & Bertaccini, Block preconditioning of real-valued iterative algorithms for complex linear systems, IMA J. Numer. Anal. 28:598-618, 2008.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
jacobi
Diagonal (Jacobi) preconditioner M^{-1} v = v / diag(A).
The cheapest useful preconditioner: one elementwise multiply per application, jit- and
vmap-native, effective on diagonally-dominant (elliptic) systems -- heat, diffusion,
elasticity. Zero/near-zero diagonals (e.g. the pressure block of a saddle-point system) are
left unscaled so it never produces inf/NaN -- but it does not rescue saddle
systems; use jno.solve.lu() or a :func:triangular block/Schur spec there.
fem.solve(linear=jno.solve.bicgstab(), precond=jno.precond.jacobi()) reproduces the
historic steady-linear default exactly.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
gmg
gmg(*, n_pre: int = 2, n_post: int = 2, omega: float | None = None, min_size: int | None = None, float32: bool = False) -> _GMG
Multigrid V-cycle preconditioner for an operator on a structured grid
(jno.shape.rect(...).structured().domain()).
The hierarchy is built from the operator itself: its per-node stencil is read by colouring
(prod(window)·nf matvecs), the coarse operators are the variational ones Pᵀ A P probed on each
coarse lattice, the smoother is ℓ¹-Jacobi with the rows' own weights (n_pre / n_post sweeps),
and the coarsest level is factorised once. Only the axes the operator couples strongly are coarsened
(semi-coarsening), and an odd cell count is merged into its neighbour, so any grid coarsens all the way.
Convergence is grid-independent on what it is given, not only on a model problem. Measured as a
standalone iteration on a 65² grid (residual factor per cycle): Poisson 0.30, a variable coefficient
0.69, a coefficient jump of 10³ 0.31, a 100:1 anisotropy 0.13, a strong reaction 0.27, an advection
term 0.57. The previous V-cycle, built from the grid alone for -Δ, gave 0.17 on Poisson and diverged
on all the others (8.3, 925, 92, 22, 1.07).
Matrix-free and differentiable. Raises if the operator has no structured grid or offers no matvec.
omega and min_size no longer apply and raise if given: ℓ¹-Jacobi has no damping parameter, and
coarsening does not stop at a size.
References: J. E. Dendy, Black box multigrid, J. Comput. Phys. 48 (1982) 366; Trottenberg, Oosterlee & Schüller, Multigrid (2001) §2.3 for the variational coarse operator; Baker, Falgout, Kolev & Yang, SIAM J. Sci. Comput. 33 (2011) 2864 for ℓ¹ smoothing.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
form
Preconditioners as weak forms: assemble an auxiliary operator  from ordinary
traced jno.fem terms and apply M^{-1} v = Â^{-1} v with inner (default
jno.solve.lu()).
This is how the classical physics-based preconditioners are written declaratively — in the same language as the PDE:
inner decides how Â^{-1} is applied. Default (None): factored once (host SuperLU)
and reused across every application — the AMS-auxiliary contract, and the only affordable
default (a per-application lu() re-factorises a 90k-DOF auxiliary hundreds of times per
solve). A jno.solve solver solves  x = v per application (an inexact inner Krylov —
pair with a flexible outer). A jno.precond spec — form([...], inner=jno.precond.ams())
— is materialized once against  and applied once per call: a multigrid or auxiliary-space
method acting on an operator declared as a weak form rather than taken from the system.
- a (weighted) mass matrix — e.g. the pressure Schur-complement approximation of a
Stokes-type saddle system:
jno.precond.form([w * pi * qi], inner=jno.solve.cg(...)); - a local proxy of a nonlocal operator — assemble only the conduction terms to precondition a conduction+radiation system (the dense view-factor coupling stays in the outer matvec);
- a shifted/damped twin of an indefinite operator (shifted-Laplacian Helmholtz);
- a low-order proxy preconditioning a high-order discretisation.
inner=False applies the assembled form instead of inverting it (M v = A v). That is
the middle factor of a product like PCD — A_p^-1 F_p M_p^-1, written
form([...]) @ form([...], inner=False) @ form([...]) — where the convection–diffusion
operator is applied between two inverses. On its own such a factor is not a preconditioner and
will make a solve worse; it exists to be composed.
The auxiliary system is assembled once with the ordinary jno.fem machinery (cached on
the spec — it is parameter-independent) and must be steady linear. Its size must match the
(sub-)operator this spec preconditions: a form over one field's symbols preconditions that
field's diagonal block inside :func:block_diag/:func:triangular; a form over all fields
preconditions the full system. A size that does not match is refused by name. On a system reduced
by a periodic tie, a slip n·u = 0 elimination or hanging-node constraints, write the form on the
full space as usual (no ties in it): it is reduced with the same field's P as the operator,
P_i^T Â P_i, when it is materialized. With an iterative inner, drive the outer solve with
jno.solve.fgmres() (flexible preconditioning).
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
inner
Use a jno.solve linear solver as the preconditioner application M^{-1} v ≈ A^{-1} v
on whatever operator it is materialized against — the natural way to give a diagonal block of
:func:block_diag/:func:triangular an (inexact) block solve, e.g.
jno.precond.inner(jno.solve.cg(tol=1e-2, maxiter=50)). With an iterative solver here the
outer Krylov must be flexible: jno.solve.fgmres().
precond preconditions that inner solve, materialized against the sub-operator this spec
is handed. Without it the inner Krylov runs unpreconditioned, which for anything harder than a
well-scaled mass matrix is no bargain — and the inner's accuracy is not a detail. Measured on a
complex A–V eddy system preconditioned through :func:real_equivalent, varying only how
accurately K + M is applied::
exact inner (sparse LU) 40 outer iterations
ONE algebraic-multigrid cycle 2200
multigrid solved to tol 1e-4 48
i.e. a loosely solved inner is worth 55x over a single cycle, and recovers essentially exact behaviour. So the useful shape is an inner solve that is cheap because its tolerance is loose, not because it is a single sweep::
jno.precond.real_equivalent(
jno.precond.inner(jno.solve.fgmres(tol=1e-4, maxiter=50),
precond=jno.precond.triangular((u, jno.precond.ams()),
(p, jno.precond.amg()))))
Do not over-tighten it either: triangular's own note records tol=1e-2 measuring 11x
slower end-to-end than 1e-4 on Taylor–Hood Stokes, the outer paying more extra iterations
than the cheaper block solve saves.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
block_diag
Block-diagonal preconditioner over the per-field DOF blocks: each (field, spec)
pair materializes spec against that field's diagonal sub-operator (field is the
trial symbol from d.fem_symbols(), or the integer block index). Cheaper per application
than :func:triangular but ignores the coupling blocks — prefer :func:triangular for
saddle systems.
On a system reduced by periodic ties, a slip elimination or hanging nodes the blocks are those of
the reduced P^T A P the solve runs on (each field's reduced DOFs), so it composes unchanged.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
triangular
Block upper-triangular preconditioner P = [[Â_1, A_12, …], [0, Â_2, …], …] over
the per-field blocks — the standard shape for saddle-point systems (Stokes / Taylor–Hood,
mixed Poisson, Biot): the last-listed block is solved first, then substituted back through
the actual off-diagonal coupling matvecs of the assembled operator.
Each (field, spec) pair supplies the approximate diagonal-block inverse Â_i^{-1}:
e.g. jno.precond.inner(jno.solve.cg(tol=1e-4)) (inexact block solve — but not too inexact:
tol=1e-2 measured 11x SLOWER end-to-end than 1e-4 on Taylor–Hood Stokes, because the
outer Krylov pays more extra iterations than the cheaper block solve saves),
jno.precond.chebyshev(...) (polynomial), or jno.precond.form([...]) (auxiliary
operator — for Stokes the classic pressure choice is the viscosity-weighted mass matrix
form([(1/mu) * pi * qi]) as the Schur-complement approximation; Elman, Silvester & Wathen,
Finite Elements and Fast Iterative Solvers, 2nd ed., OUP 2014, §9.2). With inexact
(iterative) block solves the outer Krylov must be flexible: jno.solve.fgmres().
Periodic ties, the slip n·u = 0 elimination and hanging nodes reduce the system block-wise to
P^T A P; the blocks, the couplings and any auxiliary :func:form are then those of the reduced
system, so the composition is written exactly as without them.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
lsc
Least-squares commutator Schur-complement approximation — the convection-aware one.
S^-1 ~ (B M^-1 B^T)^-1 (B M^-1 F M^-1 B^T) (B M^-1 B^T)^-1, where B is the discrete
divergence block, F the momentum (Oseen) block and M the velocity mass. Use it in place of
the pressure-mass approximation when convection matters: that recipe stands in for the Schur
complement of a VISCOUS operator, and degrades as the Reynolds number climbs.
Unlike :func:form-based approximations this is not user-writable — every factor is a block
of the assembled system rather than a weak form — which is why it is a built-in::
fem.solve(linear=jno.solve.fgmres(),
precond=jno.precond.triangular((u, jno.precond.amg()), (p, jno.precond.lsc())))
or, for the whole recipe at once, jno.precond.saddle(schur="lsc").
B and M do not depend on the solution, so P = B M^-1 B^T is assembled and factorised
once and reused across every application and every Newton step; only F varies, and it is
only ever applied. inner sets the solver for P (default an exact host LU, which keeps the
whole preconditioner a fixed linear operator). scaled=False drops the mass scaling (the
original Elman 1999 commutator) — cheaper to set up, slower to converge.
Needs an assembled operator: a steady linear solve, or Newton on the assembled tangent
(jno.solve.newton(), the default, or direct=True). A matrix-free path
(newton(direct=False)) has no blocks to slice and is refused by name. The approximation is
non-symmetric once convection is on, so drive it with gmres/fgmres, not cg/minres.
Reference: Elman, Howle, Shadid, Shuttleworth & Tuminaro, J. Comput. Phys. 227 (2008) 1790.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
pcd
Pressure convection-diffusion Schur approximation — S^-1 ~ A_p^-1 F_p M_p^-1.
The convection-aware alternative to :func:lsc, and the stronger one on the problems it suits::
fem.solve(linear=jno.solve.fgmres(),
precond=jno.precond.saddle(schur=jno.precond.pcd(viscosity=nu, inflow="inlet")))
viscosity is the diffusion coefficient of F_p and has to be given — it is physics, not
something the assembled system reveals. inflow is the coordinate tuple of the inflow
boundary, the same one the primal form used::
xin, yin, _ = d.variable("inlet", split=True)
...
jno.precond.pcd(viscosity=nu, inflow=(xin, yin))
A region name would not do: domain.variable("inlet", split=True) mints a fresh region on
every call (inlet_3, inlet_4, …), so a name re-resolved inside the spec would point at a
different region than the one the form constrained, and the Dirichlet row would quietly miss the
boundary.
Two conditions decide whether PCD works at all, both measured:
- an inflow/outflow problem. On an enclosed flow (a lid-driven cavity) PCD fails outright —
400 iterations, the solver cap, at every Reynolds number, with eigenvalues straddling zero. The
commutator argument leans on boundary behaviour an enclosed flow does not provide. Use
:func:
lscthere. - inflow-Dirichlet conditions on
A_pandF_p. With Neumann pressure operators PCD also fails (400); withinflow=given it takes 77 on the same problem and mesh. Leavinginflowunset is therefore the Neumann variant and is expected to be poor — it is allowed because it is what the literature calls the enclosed-flow form, not because it is a good default.
Measured on channel flow, momentum block solved exactly so the count isolates the Schur factor:
========== ==== ==== ==== ==== ====
Re 10 50 100 400 1000
========== ==== ==== ==== ==== ====
mass 41 184 301 284 318
lsc() 63 158 150 291 263
pcd() 37 75 77 237 355
========== ==== ==== ==== ==== ====
PCD is 4x the pressure mass at Re = 100 and loses its edge by Re = 400, where everything degrades.
Every factor is an ordinary weak form, so this spec is a convenience: the same thing is
form([...]) @ form([...], inner=False) @ form([...]), which is where to go for a variant this
does not cover (a different boundary treatment, a Robin condition, a scaled commutator).
Reference: Kay, Loghin & Wathen, SIAM J. Sci. Comput. 24 (2002) 237; Elman, Silvester & Wathen, 2nd ed., Sec. 9.2.3 for the boundary conditions.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
saddle
saddle(*, mass_weight: float = 1.0, laplace_weight: float | None = None, schur: str = 'mass', float32: bool = False) -> _Saddle
The standard saddle-point preconditioner, as one call.
Composes the classical Stokes recipe over the system's own block structure -- algebraic multigrid on the momentum block, a weighted pressure mass matrix as the Schur-complement approximation, assembled block-upper-triangularly (Elman, Silvester & Wathen, Finite Elements and Fast Iterative Solvers, 2nd ed., OUP 2014, section 9.2)::
fem.solve(linear=jno.solve.fgmres(tol=1e-10, restart=150),
precond=jno.precond.saddle(mass_weight=1.0 / mu))
which is the one-line form of::
jno.precond.triangular((u, jno.precond.amg()), (p, jno.precond.form([(1 / mu) * pi * qi])))
The constraint block is found structurally -- the field whose test function never meets its own
trial function, so the operator has no diagonal there -- which is exactly the block a diagonal
(Jacobi) preconditioner cannot touch, and the reason fem.solve()'s matrix-free default warns
on a saddle system. No symbols are passed: the pressure mass is assembled on the domain's own P1
space, so this reads the same on 2-D and 3-D Taylor-Hood.
mass_weight scales that mass matrix and is explicit, never inferred. The pressure-mass
Schur approximation wants 1/mu; recovering mu from an arbitrary weak form is fragile, and
a variable or solution-dependent viscosity would silently take the wrong weight and cost
iterations without ever failing. The default 1.0 is right for mu = 1; pass the reciprocal
viscosity otherwise. A wrong weight changes convergence SPEED only, never the answer.
Give fgmres enough restart. Its default restart=30 stagnates on a 3-D Taylor-Hood
block preconditioner, and the failure is quiet -- it still returns, just slower and less
accurate. Measured at 4302 dofs, tol=1e-10: restart=30 took 2.30 s to reach 3.3e-06,
restart=150 took 0.49 s to reach 3.3e-09 -- 4.7x faster for a thousand times the accuracy,
because restarting discards the Krylov subspace the preconditioner just earned. Larger systems
want a larger restart still.
Pair it with jno.solve.fgmres -- the pairing the block preconditioners are verified against.
jno.solve.minres does not apply: block-upper-triangular is nonsymmetric by construction and
MINRES's short recurrence assumes symmetry, so it stalls (loudly, on the residual guard, rather
than returning a plausible wrong field); use :func:block_diag if a symmetric preconditioner is
required. The non-flexible jno.solve.gmres is admissible in principle -- the mass block is
inverted exactly, so this is a fixed linear operator -- but breaks down in the default float32
build, exactly as the explicit :func:triangular composition does; prefer fgmres.
The mass matrix alone is the pure-Stokes approximation, and it degrades on a
reaction-dominated system. The pressure mass stands in for the Schur complement of
-mu*Lap(u) + grad(p); add a strong reaction term -- a Brinkman/Darcy drag alpha*u, as
porous-medium flow and fluid topology optimisation both do, or the 1/dt mass of a small
implicit time step -- and the Schur complement stops looking like a mass matrix. It still
converges, but the mesh-robustness the recipe exists for is gone.
laplace_weight is the fix: it switches the Schur approximation to Cahouet-Chabard, a
pressure mass PLUS a pressure Laplacian, S^-1 ~ mu*M_p^-1 + alpha*L_p^-1 (Cahouet & Chabard,
Int. J. Numer. Methods Fluids 8, 869-895, 1988, section 3)::
fem.solve(linear=jno.solve.fgmres(tol=1e-10, restart=150),
precond=jno.precond.saddle(mass_weight=1.0 / mu, laplace_weight=1.0 / alpha))
Both weights follow the same convention -- the reciprocal of the coefficient the term stands
for -- so mass_weight=1/mu applies mu*M_p^-1 and laplace_weight=1/alpha applies
alpha*L_p^-1. Measured through this spec on a 2-D Brinkman channel (mu = 1,
mesh_size=0.12, preconditioned GMRES to 1e-8, restart=200):
========== =========== ================
alpha mass only Cahouet-Chabard
========== =========== ================
0 83 83
1e2 135 60
1e3 503 76
1e4 2312 74
========== =========== ================
The point is not the 31x at alpha = 1e4 but the flatness: the count stops tracking
alpha. The same collapse shows with the momentum block inverted exactly rather than by AMG
(31 / 63 / 134 / 189 becoming 31 / 26 / 25 / 24 over the same four values), which is how
one can tell it is the Schur approximation doing the work and not the multigrid.
At alpha = 0 the Laplacian leg contributes nothing and the two are the same preconditioner --
the table's first row is one number measured twice. So there is no reason to pass
laplace_weight on a pure Stokes problem; it only buys an extra factorisation. Left None
(the default) the Laplacian is never assembled at all.
The pressure Laplacian is pure-Neumann and therefore singular -- the constant is an exact null vector. That is a quiet failure mode, not a loud one: a sparse LU of it factors without complaint and then applies nonsense. The auxiliary carries its own gauge pin so the factorisation is well-posed, and the applier projects the constant out on both sides to give the pseudo-inverse. There is no user knob for this and no tolerance to tune.
Scope: needs pyamg for the momentum block (a clear ImportError otherwise), and the
constraint field must be P1 -- both refuse by name rather than mis-solving. Cahouet-Chabard is
derived for a constant alpha; with a spatially varying drag (the alpha(rho) of a
topology optimisation, which is the case that most wants this) a single scalar weight is a
compromise, and the sensible choice is a representative value such as the mean over the design
field -- it degrades gracefully rather than failing, but it is no longer the sharp approximation
the table above measures. For anything else (a different momentum solver, an inexact inner solve
on a very large pressure block, more than one constraint field on mixed spaces, a genuinely
variable-coefficient Laplacian) compose :func:triangular directly; this is a shorthand for the
common cases, not a replacement for it.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
amg
amg(*, cycles: int = 1, max_levels: int = 10, coarse_size: int = 100, smoother_degree: int = 3, float32: bool = False) -> _AMG
Hybrid algebraic multigrid: smoothed-aggregation setup by the optional pyamg
(Vaněk, Mandel & Brezina, Computing 56, 1996; Bell et al., JOSS 8(87):5495, 2023), applied as
a pure-JAX V-cycle with Chebyshev polynomial smoothing (Adams et al., JCP 188, 2003) — see
:mod:jno.utils.solver.amg.
The host-side setup builds fixed-pattern level operators; the per-application V-cycle is then
jit/vmap-native and a fixed linear map, so it may precondition cg/minres as
well as bicgstab/fgmres. The mesh-independent convergence of multigrid makes this the
preconditioner for large elliptic blocks — heat, diffusion, elasticity, the (Picard-lagged)
velocity block of a saddle system inside :func:triangular.
Caching is explicit. The hierarchy is (re)built at each solve — it depends on the operator
values, so silently reusing a stale one would quietly cost iterations. To amortise the setup
over a sweep / Newton loop / inverse solve, say so: jno.precond.amg().cached(). Inside a
traced context (jit, vmap, a parametric inverse) pyamg cannot run under the trace, so build
once eagerly first — spec.build(fem.A) — and the frozen hierarchy is reused (a legitimate
preconditioner while values drift: speed degrades gracefully, correctness never). A nonlinear
fem.solve does that itself: a steady Newton builds the hierarchy once from the tangent at the
initial guess, a march from the step tangent at the initial state. pyamg is
imported lazily — without it a clear ImportError explains the install. On a matvec-only
sub-block the matrix is recovered via the (dense) block view.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
ams
AMS — the auxiliary-space Maxwell preconditioner for H(curl) (Nédélec/N1E) curl-curl systems (Hiptmair & Xu, SIAM J. Numer. Anal. 45(6):2483, 2007; Kolev & Vassilevski, J. Comput. Math. 27(5):604, 2009).
Plain point/AMG smoothing cannot damp the huge gradient near-null-space of a curl-curl
operator (curl∘grad = 0), so its condition number leaks into the iteration count. AMS adds
two corrections on cheaper nodal auxiliary problems — one on the discrete-gradient space
G, one on the vector-nodal space Π — restoring near mesh-independent convergence::
M⁻¹ r = D⁻¹ r + G (GᵀAG)⁻¹ Gᵀ r + Σ_α Π_α (Π_αᵀAΠ_α)⁻¹ Π_αᵀ r
G and Π come from the N1E edge topology (:mod:jno.utils.solver.ams); the auxiliary
operators are assembled once on the host from the concrete matrix and solved with aux.
Default (aux=None): each auxiliary block is factored once (host SuperLU) and that factor
is reused across every Krylov iteration via a host callback — the setup is amortized (the whole point
of a preconditioner) and runs off the GPU. The previous default re-factored on every iteration (a
fresh cuSolver sparse-LU per iteration — impractically slow at scale). Pass any jno.solve solver
as aux to override. Because the auxiliary problems are ordinary nodal Poisson-like systems, an
algebraic-multigrid aux makes the whole preconditioner scalable at O(n): on the GPU pass
aux=jno.solve.amg() (NVIDIA AmgX via jaxamg), which caches each auxiliary hierarchy per
operator (jaxamg.with_cache(A, is_symmetric=…)) — the AMS applier calls aux with the same
operator every iteration, so caching gives setup-once with a pure-JAX apply.
Outer solver. An exact aux (lu) is a fixed linear map, so it pairs with cg; an
inexact/iterative aux (multigrid, an inexact cg) is a variable preconditioner and
needs a flexible outer solver — :func:jno.solve.fgmres (real) — or cg stalls.
Differentiable / traced solves. The host aux-assembly cannot run under a trace. A forward
(concrete) solve freezes the auxiliaries automatically at compose time, so it is already
differentiable-ready — nothing to do. For a solve whose operator itself is traced — a jit /
vmap, or a parametric-inverse design loop where A(θ) carries a jno.np.parameter —
freeze once from a concrete reference and reuse it::
spec = jno.precond.ams().build(fem0) # fem0 = the fem at your reference parameters θ₀
node = fem_of(theta).solve(precond=spec) # parametric solve; node is differentiable
The frozen preconditioner stays valid as A(θ) drifts (speed degrades, correctness never), and
∂/∂θ flows through the operator by implicit differentiation — never through the
preconditioner (a preconditioner cannot change the solution, so differentiating its setup would be
pure waste). You cannot auto-freeze from the parametric fem itself because θ₀ is only resolved
at solve time; the one-line concrete reference is that choice made explicit.
The same spec handles the real curl-curl+mass, the complex eddy operator νK + jωσM, and
a driven time-harmonic K − k₀²εM with absorption — dtype follows the assembled matrix; pair
it with :func:jno.solve.gmres (complex-correct) for the complex cases. For a complex operator
every auxiliary is reformulated as its real-equivalent 2n block [[Re,-Im],[Im,Re]], which
a real-only aux (AmgX/multigrid) solves exactly: the gradient block GᵀA_cG and each
solenoidal block ΠᵀA_cΠ alike (non-symmetric → the aux must be non-symmetric-capable, e.g.
AMG with is_symmetric=False).
The gradient block once kept only Im A_G, on the eddy-case reasoning that GᵀKG = 0 makes
A_G = jω·R pure imaginary, so A_G⁻¹ = -j·(Im A_G)⁻¹. That fails for a driven wave
problem twice over: Re A_G = -k₀²·GᵀεMG ≠ 0, and with surface-only absorption (an impedance
/ first-order absorbing BC and no volume loss) Im A_G is a boundary mass — identically zero on
every interior node, hence singular, so the aux solve returned garbage and the outer Krylov stalled
at residual ~1 with no error. Inverting the full complex A_G fixes both, and on the eddy case
is algebraically identical to the old form (solving [[0,-R],[R,0]] gives exactly -j·R⁻¹).
Complex GMRES is not flexible, so solve a complex aux tightly (near-exact) — a strong
multigrid or lu, not a single V-cycle, which is a variable preconditioner and will stall.
Requirements & scope:
- The operator must be coercive on the gradient space — a bare curl-curl is singular there;
a mass term (conductivity, or the σ=0-in-air ε-gauge
jω·ε·⟨A,v⟩) is what makesGᵀAGinvertible. The spec raises if that term is missing. G/Πare built from the full edge topology, so this targets weak/penalty (PEC-style) boundary terms; Dirichlet-eliminated DOFs would need row-masking — out of scope here.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
jaxamg
GPU AMG preconditioner via jaxamg (NVIDIA AmgX wrapped as a JAX primitive) — the
on-device counterpart of :func:amg, and the natural smoother for large elliptic blocks or the
auxiliary nodal solves of an H(curl) AMS preconditioner on the GPU.
Builds the AMG hierarchy with jaxamg.make_preconditioner and applies a single cycle as
M⁻¹ — a proper build-once/apply-many preconditioner (unlike jno.precond.inner(jno.solve.amg()),
which re-solves each application). Wrap in .cached() to reuse the hierarchy across solves::
fem.solve(linear=jno.solve.fgmres(), precond=jno.precond.jaxamg().cached())
config is a full AmgX-format dict (default {"solver": "AMG"}). Needs an assembled
operator. Optional dependency — jaxamg (AmgX 2.5+, CUDA 12+, mpi4py/mpi4jax) is imported lazily.
symmetric=False builds a second hierarchy on A^T so the adjoint (reverse-mode) solve
is preconditioned too — an AMG hierarchy is not structurally transposable, and without this the
reverse pass of a differentiable non-symmetric solve runs effectively unpreconditioned (see
PrecondApplier). Leave the default for SPD operators, where one hierarchy serves both
directions. Real-valued operators only — AmgX has no complex mode.
Reference: Liu, Fan & Wang, arXiv:2606.09001 (2026), wrapping NVIDIA AmgX (Naumov et al., SIAM J. Sci. Comput. 37(5), 2015).
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
cached
Memoise any preconditioner's setup so it is built once and reused across solves — the plug-and-play way to amortise an expensive setup (a multigrid hierarchy, an assembled auxiliary operator, a jaxamg/AmgX coloring) over a frequency sweep, a Newton loop, or an inverse-problem optimisation, regardless of which backend does the work.
Wraps any spec (jacobi, amg, form, a jaxamg-backed preconditioner, or a user
ctx -> M⁻¹ callable). refresh=False (default) freezes the setup from the first solve and
reuses it forever — the standard frozen-preconditioner trade (a preconditioner only changes
convergence speed, never the solution, so reusing a slightly-stale setup is always correct and
usually cheap). refresh=True rebuilds when the operator's shape/sparsity changes (values may
still drift under the frozen setup); an int k rebuilds every k-th materialization — the
cadence policy for a Newton loop or transient march whose operator values drift step by step;
pass a callable ctx -> hashable for a custom invalidation key. The wrapped spec's eager prepare(fem) hook (if any) is forwarded, so it composes with the
jit/vmap/parametric-inverse build-eagerly requirement unchanged.
Only a setup built from concrete values is kept. One built inside a trace -- the per-step Newton
solve of a nonlinear march, a solve under jax.jit/jax.grad -- closes over that trace's
intermediate values, so it is used for that trace and not kept: there the cache builds once per
trace (once per compiled program), not once per run.
Reuse the SAME cached(...) object across the solves you want to share the setup::
M = jno.precond.cached(jno.precond.amg()) # backend-agnostic — pyamg or jaxamg alike
for f in freqs:
u = build_fem(f).solve(linear=jno.solve.fgmres(), precond=M) # hierarchy built once
schwarz
schwarz(*, parts: int | None = None, overlap: int = 1, coarse: bool = True, restricted: bool = False, nullspace=None, float32: bool = False) -> _Schwarz
Overlapping Schwarz domain decomposition, algebraic: built from the assembled operator alone, so it applies to any FEM/FDM system (Toselli & Widlund, Domain Decomposition Methods, 2005; restricted variant: Cai & Sarkis, SIAM J. Sci. Comput. 21(2), 1999; coarse space: Nicolaides, SIAM J. Numer. Anal. 24(2), 1987).
The unknowns are split into parts pieces by METIS on the operator's graph (default: one
part per ~256 unknowns, rounded up to a multiple of the device count), each grown by overlap layers of
neighbours; every local problem is solved exactly (the local inverses, applied as one batched product).
Unknowns with no neighbours (eliminated Dirichlet rows) are left out of the partition and solved by their
diagonal. The partition needs METIS: pip install "jax-numerical-operators[fem]" (or [fdm],
or [metis]); without it, building the preconditioner raises an ImportError saying so. coarse=True adds the two-level coarse correction (one constant per part,
applied in the balanced hybrid form), which is what keeps the iteration count from growing with the number
of parts. restricted=True is RAS -- usually fewer iterations, but non-symmetric (use bicgstab/gmres;
cg switches to flexible CG); the default is the symmetric additive form, fit for cg.
nullspace sets what the coarse space carries per part: None one constant (Nicolaides); "rigid"
the near-null space built from the problem -- rigid-body modes for a displacement-like field (as many
components as dimensions), one constant per component otherwise -- which is what elasticity needs; or an
explicit (n, k) array. On a REDUCED system (a periodic tie, slip n·u = 0, hanging nodes: the solve
runs on P^T A P) the partition is of that reduced operator, and "rigid" builds the modes on the full
mesh and restricts them to the kept DOFs -- exact for a mode the reduced space contains (a translation
along a periodic direction), a coarse vector like any other for one it does not (a rotation under a
periodic tie). An explicit array must be given on the reduced space.
Works best for elliptic, positive-definite problems (diffusion, elasticity, implicit time steps); indefinite problems (Helmholtz, saddle points as a whole) are outside what Schwarz handles robustly.
Several devices (fem.solve(shard=...), concrete or traced operator): the subdomains are split into
one block per device; each device builds, inverts and applies only its own, and their contributions are
combined by an all-reduce of an n-vector (three per application in the two-level form). The coarse matrix
and the vectors are replicated. Two caveats: the numeric setup gathers the operator's pattern keys and
values onto every device, once per factorisation (O(nnz) there, transiently -- the O(parts * m^2) local
matrices are what gets split); and a parts that is not a multiple of the device count leaves padded,
idle block slots on some devices. Verified on simulated CPU devices (placement and answers), not timed on
real multi-GPU hardware. On ONE GPU it is slower than amg.
fsai
Factored sparse approximate inverse for SPD operators: M^{-1} = G^T G ~ A^{-1} with G
lower triangular on the pattern of tril(A^power) (Kolotilina & Yeremin, SIAM J. Matrix Anal. Appl.
14(1), 1993).
The rung between jacobi and multigrid for a GPU: applying it is two sparse products -- no
triangular solves, no sequential sweep (unlike ILU / incomplete Cholesky) -- and the setup is a batch of
small independent dense Cholesky solves, one per row, so it runs in JAX, under jit and on GPU.
power=2 uses the pattern of A^2: a denser, stronger factor for a costlier setup and apply.
The pattern is computed once (host, from a concrete operator) and reused; the numeric factor follows the operator's current values, so it works for a Newton tangent or a time step whose values change. An operator whose sparsity changes after the build is detected and makes the solve fail loudly (NaN).
SPD only: a non-symmetric operator is refused; an indefinite one fails its local Cholesky (NaN). Needs x64 and an assembled operator (not a matvec-only one).
Cost, measured with CG to 1e-8 on an RTX 3070 (float64): 3-D P1 Poisson (303k unknowns) 89 iterations /
58 ms against Jacobi's 216 / 76 ms; 3-D P2 (142k) 82 / 39 ms against 197 / 58 ms; 3-D elasticity (206k)
325 / 387 ms against 778 / 588 ms. The numeric setup costs about one to two solves (101, 115, 530 ms
there), so it pays when the factor is REUSED: a constant-operator time march (built once), .cached()
across solves, or several right-hand sides. power=2 grows the setup steeply with the row width
(7.1 s on the P2 case for 32 iterations) -- try power=1 first.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
nystrom
nystrom(*, rank: int = 20, mu: float | None = None, seed: int = 0, float32: bool = False) -> _Nystrom
Randomized Nyström low-rank preconditioner for SPD operators — the rung between
jacobi and a multilevel method.
Frangella, Tropp & Udell, "Randomized Nyström Preconditioning", SIAM J. Matrix Anal. Appl.
44(2), 2023 — Algorithm 2.1 (the stabilized sketch) and §3 / Definition 3.1 (the
preconditioner P^{-1} = (lam_min+mu) U (diag(lam)+mu)^{-1} U^T + (I - U U^T)).
Sketches A against a random n x rank matrix — exactly rank matvecs, no assembled
matrix and no triangular solves — and deflates the captured top of the spectrum. That is the
part jacobi cannot reach: a diagonal preconditioner rescales, it cannot separate a few large
outlying eigenvalues, which is precisely what stalls Krylov on FEM operators with a stiff
coefficient contrast or a near-null-space. Unlike ILU it needs no factorization and no
sequential sweep, so it is jit/vmap-native and runs on GPU.
rank is the number of eigenvalues captured (cost is linear in it); mu is the
regularization, defaulting to the smallest captured eigenvalue so the low-rank and identity
parts meet continuously. seed fixes the sketch, so a solve is reproducible.
SPD only. The sketch takes a Cholesky of Omega^T A Omega, so a non-symmetric or
indefinite operator will produce NaNs rather than a wrong answer quietly — use jacobi or
chebyshev for those.
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
chebyshev
chebyshev(*, degree: int = 8, lmin: float | None = None, lmax: float | None = None, lmin_ratio: float = 1.0 / 30.0, safety: float = 1.05, bound_iters: int = 30, float32: bool = False) -> _Chebyshev
Fixed-degree Chebyshev polynomial preconditioner M^{-1} = p_degree(A) ≈ A^{-1}
for SPD operators (Saad 2003, §12.3 / Golub & Varga 1961 — the same recurrence as
jno.solve.chebyshev, truncated at degree with no convergence test, which keeps the
application a fixed linear map so it may precondition CG and MINRES).
The GPU-era substitute for Gauss-Seidel/ILU smoothing: only matvecs and AXPYs — no
reductions, no triangular solves — jit- and vmap-native.
Spectrum bounds of A are taken from lmin/lmax when given, else both ends are
measured by bound_iters steps of Lanczos (Lanczos 1950, §II — the extreme Ritz values of
the tridiagonal), at the same one-matvec-per-step cost as the power iteration it replaces.
This matters because the polynomial is a contraction only inside the interval it is fitted
to: the historical lmin = lmin_ratio * lmax guess, when it lands above the true smallest
eigenvalue, leaves the lowest modes outside that interval where the polynomial amplifies them
instead of damping. Without the optional :mod:matfree package the guess is still the
fallback (lmin_ratio then applies).
float32=True builds and applies this preconditioner in single precision while the solve stays in
double -- see :func:_materialize_in_float32.
Tracing primitives
Most users never instantiate these directly — they are what the expression-building API returns. Documented here for reference and for authors of custom operators.
jno.trace.Variable
Independent variable placeholder (e.g., x, y, t).
Carries the domain tag and dimension index so the solver can bind sampled coordinates when evaluating traced expressions.
For time-dependent problems, spatial variables (axis='spatial') index
into the spatial context array context[tag] shaped (N, D_spatial)
(after the outer B and T vmaps peel off their axes). The temporal
variable (axis='temporal') reads from a separate
context["__time__"] entry that is a scalar (after the T vmap).
jno.trace.Integral
Integral(target: 'Placeholder', integration_var: 'Variable | None' = None, quadrature: 'str | int' = 'nodal')
Mesh-based integral reduction of an expression over its domain region.
Created by :meth:Placeholder.integrate. The region (boundary vs volume)
is auto-detected at evaluation time from the Variable tags inside
target via domain._boundary_registry.
When integration_var is set (the outer/collocation Variable), the
evaluator uses jax.vmap to return an (N, 1) array instead of a
scalar, enabling non-separable Fredholm kernels.
jno.trace.Noise
Stochastic noise term regenerated every training step.
Created by :mod:jno.noise. Produces an array of shape (N, ndim)
where N is inferred at evaluation time from the number of active
spatial points and ndim (default 1) controls the trailing dimension.
The realisation is derived from the solver's step PRNG key via
jax.random.fold_in, so it is fully reproducible when the global seed
is fixed (via :func:jno.setup or .jno.toml).
Parameters
distribution : str
'gaussian', 'uniform', or 'laplace'.
**params
Distribution-specific kwargs: std, low, high, ndim.
Numerical-method front doors
Each takes a term list and returns an object carrying the assembled problem. The narrative guides are FEM, FDM and RCWA; this is the signature-level reference.
jno.fem
jno._fem.fem
Assemble a flat list of traced residuals into an :class:FEM — see :func:_fem_impl.
This thin wrapper exists only to scope chunk= over the assembly. The element-chunk policy has
to be in force while the assembler builds its closures (it is captured there, not read when they
later run), and threading a parameter through ten assemble_fem_native call sites to say the
same thing would be churn for no gain.
The FEM object
jno._fem.FEM
Assembled FEM artefacts produced by :func:fem (no solve).
Attributes
domain, mesh, problem: the owning jNO domain, its meshio mesh, and the assembled problem object. dofs: total number of degrees of freedom. classification: human-readable summary of how each residual was bucketed.
residual
property
Residual callable for a custom solver, returning a flat (n_dofs,) JAX array.
Steady linear: residual(u) is A u - b (pass args= for a runtime-parametric form).
Steady nonlinear: residual(u). Transient: residual(u, t) — the per-step semidiscrete
residual (pass args= only for a runtime-parametric solve). Use fem.operator for the raw
(unflattened) form.
Scoring a field this FEM did not produce. The residual is the only thing that can say
whether a field not obtained from fem.solve() — a reduced-basis answer, a coarse-mesh
interpolant, a neural operator's prediction — actually satisfies this system, and it costs one
sparse matvec against the O(N^3) of the solve it would certify::
rel = jnp.linalg.norm(fem.residual(u_pred)) / jnp.linalg.norm(fem.b)
That is exactly the check fem.solve(basis=...) already runs on itself (fem.basis_residual).
It is a certificate, not a bound: a small residual does not bound the error without the
operator's conditioning, and on an ill-conditioned system a plausible-looking residual can still
hide a large error. It does reliably catch a field that is simply not a solution.
The linear branch never densifies — A stays the assembled BCOO, so this is O(nnz).
jacobian
property
Jacobian callable for a custom solver, returning a dense (n_dofs, n_dofs) JAX
array — jacobian(u) (steady nonlinear) or jacobian(u, t) (transient). Use
fem.operator for the raw sparse (BCOO) form on large problems.
operator
property
The raw assembled block — (A, b) / FemLinearSystem /
FemResidualOperator (steady) or SemidiscreteTimeBlock (transient).
A COMPLEX problem returns the fused real 2n block, i.e. an ordinary (A, b) over
x = [x_r; x_i] — NOT the (re, im) leg pair it used to. Read the unfused legs off
:attr:_complex_legs instead.
Beware when migrating code: the fused value is also a 2-tuple, so op_r, op_i =
fem.operator keeps unpacking cleanly and silently binds op_i to the LOAD VECTOR. It
does not raise; it just means something else. That is what broke jno.rcwa and made a
test report a missing imaginary part that was there all along.
offsets
property
Per-field block offsets into the flat solution: sol[offsets[i]:offsets[i+1]] is field i.
Set for the native non-nodal (RT/P0) path; falls back to the problem.offset for the
Lagrange multi-field path (None for a single field with no block structure).
blocks
property
Per-field DOF slices into the flat solution (None without block structure).
The structural handle block preconditioners build on: jno.precond.block_diag /
triangular resolve their field arguments to these slices (see docs/solvers.md).
These always slice the FULL solution. When a periodic tie, slip n·u = 0 or hanging-node
constraint makes the solve run on the reduced P^T A P, a preconditioner is handed the
per-field slices of that reduced system instead (PrecondContext.blocks).
points
property
Node coordinates the DOFs live on.
For a higher-order (P2) element these are the assembly-mesh nodes
(vertices + edge midpoints), which differ from the linear mesh the
domain keeps — use these to interpret the solution vector.
field_points
property
Per-field DOF coordinates: field_points[i] is the (n_nodes_i, dim) node array the
block sol[offsets[i]:offsets[i+1]] lives on. For a coupled problem the fields may use
different nodes (e.g. Taylor-Hood P2 velocity vs P1 pressure). Backed by the per-field meshes
of the problem object, or the native assembler's recorded per-field DOF points.
stats
property
What the last :meth:solve did — observability without changing the solve's return.
None before any solve. Afterwards a dict with mode, dofs, wall_s (dispatch
time of the solve call — JAX is async; block on the result for compute time), the linear
and precond slot reprs, nonlinear (driver name, final residual norm against its
bound, step count — every driver reports it: newton_direct from its forward loop, newton_krylov and
staggered as an auxiliary output of their custom_root), and amgx_cache (AmgX solver-cache occupancy)
when jaxamg served the solve. Populated on eager paths; a solve wrapped whole in
jit/vmap/grad records the slots but no residuals — the same concrete-only
self-disabling as the convergence guards.
march is the per-step record of a march, or None: what, coord, steps,
grid (the coordinate at each step), residual / bound (each step's final residual
norm and the tolerance it was judged against — None for a march that does not judge its
steps, e.g. a linear transient) and step_s (per-step wall time — None for a march that
is one compiled lax.scan, where a step is not a host-visible event and only the mean
wall_s / steps is honest). A transient solve returns a deferred node, so its record says
deferred=True until the node is evaluated eagerly (.fn()), which fills in the
evaluation's wall_s and evaluation count. solve_index counts solves on this form:
the first includes tracing and compilation. error is set when the solve raised; the rest
of the record then describes that failed solve, not an earlier one.
is_complex
property
A complex-valued problem (steady or transient), solved via the real-equivalent block.
Both steady and transient complex problems — Bloch included — are fused into one real 2n
system at assembly, so they are ordinary "linear" / "transient" modes: a steady one
announces itself through _complex_n (the half size), a transient one through its block's
metadata["complex"]. "complex" survives only as the internal pre-fusion tag the
assembly hands to :func:_fuse_complex_steady; no finalized FEM carries it.
solve
Timed wrapper around the solve. See :meth:_solve_inner for the full signature.
eigs
eigs(*, mass, k: int = 6, which: str = 'smallest', sigma=None, linear=None, precond=None, tol=None, maxiter=None, X0=None)
Generalized eigenproblem K x = λ M x on this fem: K is this source-less fem's
operator (its stiffness bilinear form) and M is the mass bilinear form assembled on the
same space. Returns (λ, X) — the k eigenvalues at which ('smallest'/'largest')
and their M-orthonormal eigenvectors. Eigenvalues are differentiable (see :func:jno.solve.eigs).
Modal analysis, buckling, EM cavity/waveguide resonances, photonic band structure::
u, v = d.fem_symbols(); ui, vi = u.bind(x=xi, y=yi), v.bind(x=xi, y=yi)
K = jno.fem([ui.x * vi.x + ui.y * vi.y]) # stiffness (no source term)
lam, X = K.eigs(mass=[ui * vi], k=6) # K x = λ M x → λ = ω² (Neumann box)
precond= switches from the dense reduction to preconditioned LOBPCG, which never
densifies the operator and so scales past it; sigma= targets the k eigenvalues
nearest the shift (interior modes — a cavity resonance in a band, a Brillouin-zone point
away from the band edge) via the shift-invert transformation, with linear= picking the
inner solver against K − σM (default: a once-factorized host sparse LU). tol/maxiter
tune the iterative paths and are rejected on the dense one. See :func:jno.solve.eigs for the
full contract::
lam, X = K.eigs(mass=[ui * vi], k=6, precond=jno.precond.amg())
lam, X = K.eigs(mass=[ui * vi], k=4, sigma=60.0) # the 4 modes nearest λ = 60
eval
Read a term back at the solution u. What comes out depends on the KIND of expression.
A weak term (one carrying the test function) assembles to the free (n_dofs,) vector, with
no essential elimination applied. An expression with no test function is an integrand — an
energy or stress density — and reduces instead to the scalar ∫ F dΩ over its region, which is
how an objective (compliance, a volume fraction, a stress p-norm, a mechanism output) is written.
The two are described in turn below; a list mixing them is refused, since they have no common
return shape.
The weak term: a free residual.
This is the readout primitive. The conjugate quantity on a constrained region — reaction force in mechanics, total heat flux through a Dirichlet wall, current in electrostatics, flow rate in Darcy — is that vector summed over the region's DOFs::
fem = jno.fem([mech, *bcs])
u = fem.solve()
R = fem.eval(mech, u) # internal force at every DOF
Fy = R[fem.region_dofs("top", component=1)].sum() # reaction on the top edge
It has to be a separate entry point because every solve path elimination-mutates the system it
keeps: the linear route applies symmetric elimination (A/b here have the constrained rows
zeroed and a unit diagonal set), and Newton replaces those rows with u[d] - g. Both are
correct for solving and both return exactly zero at the DOFs a reaction readout is asking
about — so reading it off fem.A/fem.b/fem.residual gives a plausible, silent zero.
term is any weak term built from this domain's symbols; it does not have to be one of the
terms this FEM was built from, so a diagnostic form — a sub-term, a different stress measure — can
be assembled against an existing solution. On a transient problem a weak term is assembled at
t = 0.
The integrand: a domain integral. Drop the test function and the same call returns a scalar::
C = fem.eval(sig(u, rho) @ eps(u), sol, args=design) # ∫ σ(u):ε(u) dΩ — the compliance
V = fem.eval(rho, sol, args=design) # ∫ ρ dΩ — the material volume
g = fem.eval(inner(u.bind(x=xo, y=yo), e), sol) # ∮ u·e ds — a mechanism's output
It integrates on the quadrature the operator was assembled with, so
fem.eval(a(u,u), sol) == sol @ fem.eval(a(u,phi), sol) exactly — the same rule on both sides,
rather than a difference nothing would report. Differentiable in the solution, in the runtime
parameters, and in the mesh coordinates (∂/∂X flows through |det J| and the facet area
element), which is what a deformable-mesh design problem needs.
The region is read off the expression's coordinate Variables exactly as for a weak term, so a
boundary integrand is spelled by binding the field to that region's coordinates; a term naming two
regions is refused, because an integral has one measure. Narrower in scope than the weak path: the
whole volume and tagged boundary regions only, and steady native-Lagrange problems only. Where it
is unavailable, sum(fem.eval(F * phi, u)) is the same integral — a Lagrange basis is a
partition of unity.
block_index
Resolve a trial symbol (or plain index) to its position in :attr:blocks /
:attr:offsets — the field order is first appearance in the jno.fem constraints.
region_dofs
Global DOF indices of a tagged region — the companion to :meth:eval.
field selects the block (a trial symbol or an index, resolved by :meth:block_index);
component picks one component of a vector field (None = all of them). Returns a plain
numpy int array, so it indexes a solution or a residual directly.
jno.fdm
jno.fdm.fdm
Finite-difference PDE solver — the strong-form sibling of :func:jno.fem.
The problem is a constraint list, exactly as in jno.fem: jno.fdm([residual, u(xb, yb) - g])
authored with u = domain.unknown() and strong-form derivatives. Constraints are classified by
region — interior terms are the PDE residual, boundary terms are Dirichlet conditions — and a
u.t term makes it transient, taking its grid from domain.time.
For the accurate whole-Laplacian stencil write ONE term,
u.laplacian(x, y, scheme="finite_difference:cotangent"); per-axis d2 refuses that
sub-scheme, because summing one per axis would multiply the Laplacian by the dimension.
jno.fd
jno.stencils.fd
fd(*, order: int | None = None, points: tuple[int, ...] | None = None, weights: dict[int, float] | None = None, boundary: int | None = None, average: str | None = None, upwind: Any = None, fit: int | None = None, rings: int | None = None) -> FDStencil
A finite-difference stencil, for scheme=: see the module docstring.
| PARAMETER | DESCRIPTION |
|---|---|
order
|
accuracy order of a central stencil (even). Default 2.
TYPE:
|
points
|
integer offsets the stencil reads, e.g.
TYPE:
|
weights
|
TYPE:
|
boundary
|
order of the one-sided stencil used where the interior one does not fit. Default: the interior order.
TYPE:
|
upwind
|
the wind (an expression: a number, a formula, a field, the unknown itself; a vector wind gives
its component along each axis). The first derivative is then upwind-biased per node, of order
TYPE:
|
fit
|
unstructured meshes: the degree of the local polynomial least-squares fit the derivatives are read
from (default
TYPE:
|
rings
|
unstructured meshes: how many rings of neighbours the fit uses (default: the fewest giving every node 1.5× as many neighbours as the fit has coefficients).
TYPE:
|
average
|
the coefficient between nodes in the conservative form of
TYPE:
|
jno.rcwa
jno.rcwa.rcwa
rcwa(problem, *, orders, wavelength=None, grid=64, nz=64, slices=None, params=None, formulation='JONES_DIRECT_FOURIER', smoothing=1, polarization=None)
Infer and build an RCWA problem from a jNO constraint list (or built FEM).
Everything is read out of the traced problem: periodicity + period (Floquet ties — absent ⇒
raise), the super/substrate ambients, the permittivity (the K0**2*eps coefficient
recovered from the scalar Helmholtz volume term, sampled along z; a tensor ε̂ from
inner(ε̂ @ u, v)), the wavelength (k0 from the vacuum superstrate, unless wavelength
is given), and the incident wave (illuminated face + transverse angle k_in, from the
assembled forcing). An in-plane PML (a complex coordinate stretch in the stiffness coefficients)
is detected and honoured as a diagonal Maxwell ε̂/μ̂ — the supercell then behaves as an
isolated scatterer, and a design parameter on the scatterer ε is differentiable (inverse design of
an isolated structure). Scope: uniaxial (diagonal) in-plane stretch only, scalar Helmholtz.
An internal source (a - f·v / - inner(J, v) volume forcing) switches the solve to an
emission problem: the source is localized (point vs Gaussian; dipole orientation from J), the stack
is split at the source plane, and sol.power("up"|"down"|"total") / sol.extraction("up"|"down")
give the radiated power and directionality. The amplitude, orientation, lateral location, z-position and
Gaussian width, plus the design ε, are all differentiable (only the discrete layer choice is frozen);
k_in is nudged off the singular Γ-point, and the source must sit in a finite (contrast-defined) layer.
Scalar Helmholtz only for now; a boundary incidence + internal source together raise.
Parameters
problem:
The same constraint list you would pass to :func:jno.fem, or an already-built FEM.
orders:
Fourier truncation for the modal solve.
wavelength:
Free-space wavelength (same length unit as the geometry). Optional — inferred from the
superstrate (assumed vacuum) when omitted. Pass it to override or if no ambient is vacuum.
grid, nz:
Transverse resolution and number of z-samples used to detect the layer stack.
slices:
Staircase a continuously-varying permittivity into this many layers instead of raising.
smoothing:
Subpixel-smoothing factor (default 1 = off). smoothing=k supersamples each RCWA pixel
k×k and area-averages the permittivity, anti-aliasing material boundaries. Recommended (2–4)
for inverse design: it cuts Fourier (Gibbs) ringing so fewer orders converge, and makes the
gradient w.r.t. a boundary-moving design parameter smooth rather than staircased. Costs k²× more
coefficient evaluations (cheap) — the fmmax eigensolve size is unchanged.
polarization:
The incident polarization: "x", "y", or a transverse Jones vector (p_x, p_y) (complex
entries give elliptical light). Usually inferred, and the choice is logged and kept on
spec.polarization / spec.polarization_from. A vector (curl-curl) list carries it in its
incident field E_inc. A scalar Helmholtz list is solved in the polarization for which Maxwell
reduces exactly to the written equation -- E along the axis every layer is invariant along, when the
incidence has no component along it (a 1-D grating: E parallel to its lines; TE) -- so jno.rcwa
and jno.fem answer the same equation. A stack patterned along both axes has no such
polarization; it is solved x-polarized, and the log says so.
Returns
_RcwaProblem
Holds the inferred :class:RcwaSpec (.spec) and a :meth:~_RcwaProblem.solve.
jno.rcwa.Rcwa
Rcwa(layers, *, period, orders, wavelength=None, k_in=(0.0, 0.0), formulation='JONES_DIRECT_FOURIER', assume_periodic=False, polarization='x')
A periodic layered RCWA problem with an explicit layer stack (the backend engine).
Most users reach this through :func:jno.rcwa with a problem list; construct it directly only when
you already have the [(thickness, eps), ...] stack, period and truncation in hand.
solve
Solve the stack and return a :class:_Sol. Raises if the wavelength is unknown or energy is
not conserved.
layers optionally overrides the construction layer stack with [(thickness, eps), ...] at
solve time -- pass JAX permittivity grids here to differentiate the solve in the design (the
construction-time shape guards run once, eagerly, so the solve itself stays trace-clean).
source (a dict built by the front door) drives an internal-source emission solve instead of
plane-wave incidence -- see :meth:_solve_source; it returns an :class:_EmitterSol.
jno.rcwa.RcwaSpec
dataclass
RcwaSpec(period: tuple, layers: list, wavelength: float, k_in: tuple = (0.0, 0.0), source_face: str = '', ambient_faces: tuple = (), periodic_axes: dict = dict(), source: dict = None, polarization: tuple = (1.0 + 0j, 0j), polarization_from: str = '')
Everything the RCWA engine needs, inferred from a jNO problem. Inspectable without fmmax.
The RCWA solution
Returned by rcwa(...).solve(). Every readout is a differentiable JAX array, so a transmission or
per-order objective can be optimised straight through the modal solve.
jno.rcwa._Sol
efficiency
Total transmitted ("T") or reflected ("R") power fraction over propagating orders.
Returns a JAX scalar -- differentiable in the design (permittivity) so jax.grad of a
transmission/reflection objective flows through the modal solve.
order
Diffraction efficiency into transmitted order (m, n) (raises if outside the truncation).
Returns a differentiable JAX scalar.
jones
The 2×2 complex Jones matrix at the 0th diffraction order — how the structure maps incident
polarization to transmitted ("T") or reflected ("R") polarization.
J[q, p] is the 0th-order field amplitude in output polarization q for a unit-power
incident wave in input polarization p, normalised so |J[q, p]|² is the (co- or cross-
polarised) power fraction and arg(J) carries the phase (a waveplate's retardation lives in the
phase difference between the two diagonal entries). The columns sum in power to the total efficiency:
sum_q |J[q, p]|² == efficiency(kind) for input p. The two polarizations are fmmax's transverse
Jones basis — for the 0th order at normal incidence they are the two in-plane axes (≈ x, y); an
isotropic stack gives a diagonal J (no conversion), an in-plane-anisotropic one an off-diagonal
J (polarization conversion). Differentiable in the design. Returns a JAX (2, 2) array.
field
Reconstruct the real-space electric field on a vertical (x–z) slice at y = y_frac·Py.
Returns (intensity, extent, layer_z): intensity = |E|² of shape (nz, nx) (z
vertical, x horizontal) for imshow; extent = [0, Px, z_min, z_max]; layer_z = the
z-interfaces between layers (to annotate the patterned slab). Needs the finite ambient
thickness recorded at solve time (semi-infinite ambients are shown as unit-thick slabs).
See :meth:field3d for the whole volume rather than one slice.
field3d
|E|² on the full 3-D grid, for volume rendering / isosurfaces — the whole reconstruction
rather than the single y-slice :meth:field returns.
Returns (intensity, (x, y, z), layer_z) with intensity of shape (nx, ny, nz) and
x/y/z the 1-D coordinates of each axis. The semi-infinite ambients are rendered as
unit-thick slabs (see :meth:field), so z starts at 0 at the top of the superstrate, not at
any coordinate of the authored geometry — layer_z gives the interfaces in the same frame.
NumPy, like :meth:field: this is the visualisation path. A differentiable |E|² for a
coupled (e.g. thermal) objective needs the JAX-valued volume, which this deliberately does not
return — see the roadmap in docs/rcwa.md.
aerial
Partially-coherent aerial image (wafer-plane intensity) of this mask, by Abbe source integration -- the litho imaging step on top of the rigorous RCWA mask diffraction.
The mask's diffraction orders (this solution's spectrum) are projected through a lens of numerical
aperture NA, summed over the illumination source. Everything (wavelength, period, the complex
mask spectrum) is read from the solve; only the optics are yours.
Parameters
NA:
Numerical aperture of the projection lens (an order passes the pupil when its direction cosine
|α| ≤ NA).
source:
Illumination. A float σ → conventional (a filled disk of partial-coherence radius σ);
a (σ_in, σ_out) tuple → annular; a raw array of pupil weights (differentiable, for
source-mask optimization).
defocus:
Defocus (same length unit as the geometry); applies the quadratic pupil phase.
grid:
Real-space resolution of the returned image (one period).
kind:
"T" (transmissive DUV mask, default) or "R" (reflective EUV mask).
polarization:
None (default) → scalar imaging of the field component along the incident polarization
(E_x for an x-polarized solve), correct at low NA. A string turns on the
vector high-NA model, which rotates each order's transverse (E_x, E_y) to the 3-D wafer
field through the Richards-Wolf/Flagello vector pupil (so the TM component loses contrast at large
angles): "x" / "y" → linearly polarized illumination; "unpolarized" → the two averaged.
Returns
A (grid, grid) JAX intensity array over one period -- differentiable in the mask design and the
source, so jax.grad of a printed-vs-target loss drives OPC / ILT / SMO.
Notes
Uses the shift-invariant (Kirchhoff-like) Abbe sum on this single solve's rigorous spectrum; full
angular rigor (re-solving the mask per source point) is a heavier option. The vector pupil follows
Flagello, Milster & Rosenbluth, J. Opt. Soc. Am. A 13, 53 (1996), with the aplanatic
1/√(cosθ) apodization; at NA→0 it reduces to the scalar image.
jno.rcwa.RcwaError
Raised for any ill-posed RCWA problem — always loud, never a silent wrong answer.
Geometry
jno.geometry.shape.shape
dataclass
shape(_node: tuple, dim: int, _size: Size = None, _region_name: Optional[str] = None, _mesh_order: int = 1, _attach: Optional[dict] = None, _cell: Optional[str] = None, _structured: Optional[tuple] = None)
An immutable geometry build-plan. Operators return new shapes.
Derive with :func:dataclasses.replace, never by calling shape(...) positionally. Every
method here returns a copy with one or two fields changed, and a positional constructor has to
re-list every other field to carry it — so a field added later is silently dropped by whichever
derivation forgets it, and the plan quietly meshes as something the caller did not ask for. That
is not hypothetical: .quad().attach(k=1.0) used to erase the quadrilateral choice outright.
polygon
classmethod
Arbitrary 2-D polygon from ordered (x, y) vertices; edges auto-named e0, e1, ....
cylinder
classmethod
cylinder(x: float, y: float, z: float, dx: float, dy: float, dz: float, r: float, size: Size = None) -> 'shape'
Right cylinder: base centre (x,y,z), axis vector (dx,dy,dz), radius r.
sphere
classmethod
Sphere centred (cx,cy,cz) radius r.
regions
classmethod
A multi-material domain: named sub-regions meshed conforming to their interfaces.
Each keyword names a sub-region shape (all same dim); the pieces are fragmented
(occ.fragment) so element edges align exactly with every material interface, then each
volume cell is assigned to the first region (keyword order = priority) whose shape
contains its centroid — exact, because the mesh conforms. The realized domain exposes each
region as its own variable set (d.variable("core")) alongside interior/boundary,
and the outer boundary keeps its auto-names; internal interface facets are not boundary.
Regions may overlap: shape.regions(inclusion=disk, matrix=plate) labels the disk
inclusion (higher priority) and the remainder matrix. Equivalent to combining named
shapes with + — disk.name("inclusion") + plate.name("matrix"). Must be the top-level
shape (call .domain() on it); it is not composable with boolean operators/transforms.
conforming=False skips the fragment: every piece is meshed independently, so two
touching regions end up with two coincident but non-matching surfaces and duplicated nodes
instead of one shared interface. Each side is tagged "a|b.a" / "a|b.b", and tying them
with u("a|b.a") - u("a|b.b") in jno.fem glues them by a mortar coupling — which is how
you join two bodies meshed at different resolutions. conforming is therefore a reserved
region name.
A positional {name: shape} dict is accepted for names that are not valid Python
identifiers -- shape.regions({"Quartz.1": q1, "Quartz.2": q2}). Dict order is priority
order, exactly as for keywords, and the two forms may be combined (dict entries first).
name
Label this shape as a named material region, for combining with + (see :meth:regions).
core.name("core") + clad.name("clad") builds a multi-material domain whose regions are
core and clad. Apply name last (a later transform drops the label).
attach
Attach material properties to this region: .attach(k=220.0, eps=0.794).
The realized domain exposes each attached name as a per-region coefficient ready to drop
into a weak form -- d.k is exactly d._by_region({"Kristall": 220.0, ...}) assembled from
every region that attached a k::
kri = shape.polygon(v).name("Kristall").attach(k=220.0, eps=0.794)
gas = shape.polygon(w).name("Gas").attach(k=0.186, eps=1.0)
d = (kri + gas).domain()
heat = d.k * (T.x*s.x + T.y*s.y) - d.q * s
A value may be anything :meth:domain.attach accepts -- a scalar, a symbolic expression,
or a traced/trainable array -- so an attached property can be fitted or differentiated
through. A plain function is also accepted and is called with the domain's spatial
coordinates when the property is read (.attach(k=lambda r, z: 2.0 + 0.5*z)); it has to be
deferred that way because a spatially varying coefficient is built from d.variable(...),
which does not exist yet while the geometry plan is being written.
Properties are declared, not typed: whether eps is a volume or a surface quantity is
decided by the term that consumes it, not here.
d.<name> raises if any region failed to attach that name, listing the ones that did
not: a forgotten material surfaces at first use rather than as a region that silently
conducts nothing. Repeated calls merge (last wins), so properties can be built up in stages.
Apply after :meth:name -- like name, a later transform drops the attachment.
size
Alias for :meth:sized -- .size(h) reads better in a chain than .sized(h).
revolve
Sweep a 2-D shape around an axis by angle radians into a 3-D solid.
angle == 2*pi gives a full solid of revolution; a partial angle gives a wedge
or half-donut (revolve a disk offset from the axis by pi). The profile lies in
the z=0 plane on the positive-radius side; the axis must be the x- or y-axis through
the origin (the common axisymmetric case) -- other axes raise at build time. Swept
faces inherit the profile edge's auto-name (arc, e0, ...); a partial sweep
adds back/front end caps.
translate
Move the shape by vector (2- or 3-component). Boundary names are preserved.
rotate
Rotate angle radians about the axis through axis_point along axis_dir.
sweep
Sweep this 2-D profile along an open :class:~jno.geometry.path.Path trajectory.
The path must be smooth (line and arc segments); a sharp line->line corner is rejected
up front (it self-intersects the swept profile). Naming degrades to interior +
boundary for a general sweep -- carve regions afterwards with d.tag(predicate).
array
n fused copies of this shape: a linear array (step= vector between copies)
or a polar array (about=(axis_point, axis_dir) spread over angle).
Pure composition over translate/rotate/fuse -- e.g. a bolt-circle of holes is
plate - shape.disk(R, 0, r).array(8, about=((0,0,0),(0,0,1))).
fillet
Round the solid's edges by radius.
where=f(x, y, z) selects which edges to round by their midpoint (default: all
edges). The rounded blend faces are unnamed (they fall into boundary); the flat
faces keep their names.
sized
Return a copy of this shape with its target mesh size set (scalar or f(x,y,z)).
curved
Return a copy meshed with curved (isoparametric) geometry of the given order.
By default jNO meshes straight-sided and synthesises any higher-order nodes at the
straight-edge midpoints, so the domain stays a polygon however high the element order goes.
That polygonal approximation carries an O(h²) domain error at every basis order — it is
what caps P2/P3 at second order on a round boundary — and leaves facet normals O(h) wrong.
curved() asks the CAD kernel to place those nodes on the true surface instead::
d = jno.shape.disk(0, 0, 1, size=0.1).curved().domain()
Worth pairing with a matching element order (jno.fem(..., order=2)); a curved mesh under a
P1 basis buys only the geometry, not the convergence rate. Meshing is a property of the shape,
which is why this lives here beside :meth:sized rather than on the solve.
quad
Return a copy meshed with quadrilaterals — or, on a structured 3-D plan, hexahedra::
d = jno.shape.disk(0, 0, 1, size=0.1).quad().domain() # quads
d = jno.shape.box(0, 0, 0, 1, 1, 1, size=0.1).structured().quad().domain() # hexes
In 2-D gmsh meshes triangles and then recombines them, which works on arbitrary geometry — a disk recombines to pure quads just as a rectangle does. Quadrilaterals cost fewer cells for the same nodes and are the element the topology-optimisation literature is written on; on bending-dominated elasticity they are markedly less stiff than linear triangles.
In 3-D it needs :meth:structured. gmsh cannot hex-mesh general geometry: measured here,
Recombine3DAll on a plain box returns 944 tetrahedra and no hexahedra at all. Hexes come
from a regular lattice (or, later, from sweeping/transfinite meshing on a swept build plan),
so an unstructured 3-D shape is refused and the message says to add .structured().
Meshing is a property of the shape, which is why this lives beside :meth:curved and
:meth:sized rather than on the solve.
tri
Return a copy meshed with simplices — triangles in 2-D, tetrahedra in 3-D.
Simplices are the default, so this is the explicit opposite of :meth:quad: it turns
recombination back off for a shape that inherited it from an enclosing plan::
plate.quad() - hole.tri() # (once mixed-cell meshes land, see below)
(plate.quad()).tri() # today: cancels the .quad()
A cell choice rides on the shape it describes, so each shape in a plan already carries its
own. What is not built yet is honouring two different choices in one mesh: gmsh can do it
(setRecombine is per-surface, and a quad region and a triangle region meeting at a shared
edge conform node-for-node, since both have 2-node edges there), but jNO's assembler is built
around one element table and one cell array. Until that lands, a plan that asks for two
different cells is refused rather than silently meshed with one of them.
structured
Return a copy meshed as a regular lattice instead of by gmsh::
jno.shape.rect(0, 0, 1, 1, size=0.1).structured().domain()
jno.shape.box(0, 0, 0, 1, 1, 1, size=0.1).structured().quad().domain() # hexes
jno.shape.box(...).structured(n=(32, 16, 16)).domain()
Three things follow from a lattice that a gmsh mesh cannot give:
- Hexahedra. gmsh cannot hex-mesh general geometry, which is why :meth:
quadrefuses a 3-D shape — but a structured box is exactly the plan that can be hex-meshed, so.structured().quad()is how a hex mesh is spelled. - A grid descriptor on
domain.grid, which is what letsjno.fdmtake its assembly-free 5-/7-point stencils instead of the cotangent operator, and what a nodal field reshapes against for operator-learning work. - Exactly matched opposite faces, so whole-domain periodic ties collapse onto one DOF rather than holding to a tolerance.
n counts cells, so a lattice has n + 1 nodes per axis and domain.grid["shape"]
is n + 1 — consistent with the nx/ny/nz of every other grid in jNO. Pass a
scalar for every axis, a tuple per axis, or nothing at all to derive it from the shape's own
size= (n = max(2, round(extent / h))), which keeps one resolution concept.
Refuses by name -- a CSG plan, a non-rectangular primitive, a graded size= callable --
rather than falling back to gmsh: a caller who then reads domain.grid or expects hexes
would fail somewhere else instead, having silently solved on a different mesh.
Meshing is a property of the shape, which is why this lives beside :meth:quad and
:meth:curved rather than on the domain -- by the time a domain exists it holds a mesh and
everything derived from it, and a mesh loaded from a file has no geometry to rebuild from.
cell_choices
Every explicit cell choice (:meth:quad / :meth:tri) anywhere in this build plan.
More than one means a mixed-cell mesh was asked for. gmsh can produce one — recombination
is per-surface, and in 2-D a quad region and a triangle region sharing an edge conform
node-for-node because both have 2-node edges there — but jNO's assembler is built around a
single element table and a single cell array, so the mesh could be built and not assembled.
:func:emit.build refuses on that rather than quietly meshing everything one way.
leaves
Flat (primitive, size, key) list of every primitive in the plan.
A size set on a compound shape ((a - b).sized(h), core.name("core").sized(h))
is inherited by every primitive leaf underneath it that has no size of its own — so
per-region .sized() on a CSG region actually reaches the mesher's per-primitive
Distance/Threshold fields instead of being silently dropped. A leaf's own size wins.
contains
Boolean mask over points (shape (N, dim)): which lie inside this shape.
Analytic CSG membership — primitive leaves combined by -
(cut), | (fuse), & (inter) — so it needs no gmsh and works in 2-D and 3-D alike. This is
the shapely-free point-in-region test that resolves a geometric domain.region(name, shape) to
a mesh-node subset for a subdomain / domain-decomposition solve. Inclusive within tol (the
analogue of shapely's buffer(1e-9)), so nodes exactly on a face count as inside.
translate/rotate/extrude/revolve are handled by mapping the query point into
the child's frame; sweep and fillet have no closed form and raise
:class:NotImplementedError (sample those through :meth:sample_interior, which falls back
to the boundary tessellation).
sdf
Signed distance to this shape: negative inside, zero on the boundary, positive outside.
Two calling forms, because the two places it is useful speak different languages::
shape.sdf(points) # (N, dim) numpy or jax array -> an array
shape.sdf(x, y) # jno Variables -> a trace node
The array form backs sampling, membership and mesh sizing. The trace form is the one that
composes into a weak form or a PINN ansatz, where the coordinates are Variable nodes and
arithmetic has to be emitted through jno.np rather than evaluated. The same per-primitive
formulas serve both -- only the array module differs.
Two things it is for::
u = g + shape.sdf(x, y) * net(x, y) # a boundary condition that HOLDS, not one
# that is penalised: the factor is exactly
# zero on the boundary by construction
eps = 1 + (eps_r - 1) * 0.5 * (1 - jno.np.tanh(shape.sdf(x, y) / w))
# a material that varies by region without a
# conforming multi-material mesh
contains is the sign of this: contains(p) == (sdf(p) <= 0).
Exactness. Per-primitive distances are exact, and rigid transforms preserve them. Boolean
composition (min/max) is exact ON THE BOUNDARY -- the zero set is right, which is all a
hard boundary condition or a smoothed indicator needs -- but is only a bound strictly inside a
composite. Do not lean on the interior value where the true distance matters.
Differentiability. Differentiable in the query coordinates, with a finite gradient everywhere -- inside, on the boundary and at a centre, in float32 as in float64. Not in the shape's own parameters: radii, corners and vertices are static floats read at trace time, so optimising the geometry itself does not go through here.
sweep/fillet have no closed form and raise; revolve is not wired yet.
bounds
Axis-aligned bounding box as (lo, hi), two 3-vectors — no gmsh.
Always a superset of the true extent, which is what makes it a valid proposal box for
the rejection sampler in :meth:sample_interior. Raises for a sweep/fillet plan,
whose extent is only known from its tessellation.
is_analytic
Whether this plan can be sampled with no gmsh at all.
All three of extent, membership and a boundary sampler have to be closed-form: a shape that knows where it is but cannot produce a point on its own surface is no use to a PINN, which needs boundary collocation as much as interior. A plan that fails any of them keeps the eager meshing path rather than being half-served.
tessellate
Mesh this shape's BOUNDARY only -- perimeter in 2-D, surface in 3-D -- and cache it.
The volume fill never runs, so this is far cheaper than :meth:build, and the result answers
the three questions a fillet plan cannot answer for itself: where its surface is, which
way that surface faces, and what lies inside it. :meth:contains,
:meth:sample_interior and :meth:sample_boundary fall back to it automatically; call it
directly only to inspect the facets.
Cached on the shape, so the mesher runs once. See :class:BoundaryMesh for the accuracy this
buys and gives up -- exact at creases, O(h) on curved boundaries.
sample_interior
(n, 3) points drawn uniformly from the interior — the mesh-free collocation draw.
Rejection sampling in the analytic bounding box, so points are continuous: they are not
drawn from any fixed node set and two calls give different points. A 1-D shape has no
volume to reject into and is parametrised along its arclength instead (see
:meth:jno.geometry.primitives.Curve.sample_interior).
Falls back to the boundary tessellation for a sweep/fillet plan; that path logs
what it did rather than pretending it stayed analytic.
sample_boundary
(points (n,3), normals (n,3)) uniform by measure on the boundary, normals outward.
Points land exactly on the analytic boundary — on the true circle, not on a chord — and the normal is the closed-form one, so a disk's normal is exactly radial rather than piecewise-constant per facet.
Candidates come from each primitive's own boundary, weighted by measure; a candidate
survives only if the shape actually changes across it (contains differs a hair either
side), which is what trims the parts of a leaf's boundary that a boolean cut away and what
orients the normal outward — no winding convention is assumed.
edge
Select boundary entities carrying auto-name name ("top", "left", ...).
edges_from
Select boundary entities that originate from the primitive(s) in sub.
build
Mesh this shape -> (meshio.Mesh, dim, ds). Imports gmsh lazily.
algorithm selects gmsh's meshing kernel and threads its thread count. Which kernel
it applies to follows from this shape's own dimension, so there is no separate 2-D/3-D
argument. None keeps jNO's defaults -- see :mod:jno.geometry.emit, which records why
each was chosen; the 3-D default is HXT on 8 threads, measured 14.5x faster than gmsh's
serial Delaunay at equal-or-better element quality.
domain
Build a jno.domain from this shape as a one-liner.
Forwards the domain keyword arguments (time=, sample=, name=, ...) to
jno.domain -- but not a constructor or mesh_size: the mesh size lives on the
shape itself (via size= / :meth:sized). So shape.rect(0, 0, 1, 1, size=0.1).domain()
replaces jno.domain(shape.rect(0, 0, 1, 1, size=0.1)), and batching still composes as
B * shape.rect(...).domain().
jno.geometry.path.Path
A contour/trajectory under construction. Start at a point, then chain segments.
line_to
Straight segment to (x, y, z).
arc_to
Circular arc to (x, y, z) passing through the point through (a 3-point arc).
face
Close the contour (using its x,y) and return it as a 2-D :class:~jno.geometry.shape.shape.
curve
Return this open path as a 1-D :class:~jno.geometry.shape.shape (a curve domain).
The 1-D sibling of :meth:face: face closes the contour into a 2-D region, curve
keeps it open as a 1-D manifold. The two overall endpoints are named left (start) and
right (final end) -- so jno.Path(0, 0).line_to(1, 0).curve(size=0.01).domain() is a
unit interval whose ends you reach via d.variable("left"/"right"/"boundary"). Segments
keep their 3-D points, so arc_to gives a curved 1-D manifold and multiple segments a
polyline (intermediate junctions are interior).
Optimizers (jno.optimizers)
Optax-compatible transformations. Anything optax exposes works too — these are the additions jNO needs for PDE-constrained and topology-optimisation work.
jno.optimizers
jno.optimizers — custom optax-compatible optimizers that aren't in optax.
Just the optimizers — for everything else (chain, learning-rate schedules, gradient clipping, …)
use optax directly. Each optimizer here is an optax GradientTransformation, so it composes
with optax.chain and drops straight into model.optimizer(...) /
jno.np.parameter(...).optimizer(...)::
import optax, jno
k.optimizer(jno.optimizers.ssbroyden()) # use it directly
k.optimizer(optax.chain(optax.clip_by_global_norm(1.0), # …or inside an optax chain
jno.optimizers.ssbfgs()))
- :func:
ssbroyden/ :func:ssbfgs— Self-Scaled Broyden / BFGS quasi-Newton with a line search (Urbán, Stefanou & Pons, J. Comput. Phys. 523 (2025) 113656); excellent on smooth PINN / inverse losses. - :func:
soap— SOAP, Shampoo with Adam in the preconditioner's eigenbasis (Vyas et al. 2024, arXiv:2409.11321). - :func:
mma— Method of Moving Asymptotes (Svanberg, IJNME 24(2), 1987), the standard optimiser for constrained structural design. Unlike everything else here it is a sentinel, not aGradientTransformation: an optaxupdate()never sees the individual constraint values and gradients MMA needs. Pair it with :func:jno.le/ :func:jno.ge. - :func:
md— Magnitude–Direction Decoupling: a generic wrapper that factorizes each weight matrix into a fixed-norm direction + learnable magnitude gains and steps any optax base optimizer on the direction (Hägele et al. 2026, arXiv:2606.25971). Pass it as a sentinel vianet.optimizer(jno.optimizers.md(optax.adam(1e-3))); :func:md_decoupleis the bare transform.
The scale_by_* variants are the bare transforms (no line search / learning-rate wrapper) for
building your own optax chains.
Add a custom optimizer: drop an optax-compatible GradientTransformation into a new module in
this package and re-export it below.
GeometricDecay
Multiply a scalar parameter by gamma every iteration. See :func:geometric_decay.
HeavisideContinuation
HeavisideContinuation(param: Any, *, start: float = 1.0, maximum: float = 16.0, over: int, hold: int = 0, schedule: str = 'geometric')
Ramp a projection sharpness from start to maximum. See :func:heaviside_continuation.
SIMPContinuation
SIMPContinuation(penal: Any, density: Any, *, start: float = 3.0, step: float = 1.0, maximum: float = 8.0, tol: float = 0.0001, window: int = 3, mnd_tol: float = 0.0001, physical: Optional[Callable] = None, watch: Optional[int] = None, every: int = 1, patience: Optional[int] = None)
Raise the SIMP exponent on the paper's schedule. See :func:simp_continuation.
ENGDOptimizer
Sentinel stored on a model via net.optimizer(jno.optimizers.engd(...)).
Not an optax GradientTransformation — jno.core.solve() detects
this object early, auto-detects gram_terms from all loss expressions
involving the model, injects optax.sgd(1.0) as the actual parameter
update transform, and prepends an
:class:~jno.utils.adaptive.callbacks.ENGDCallback.
See :func:engd for the public constructor.
MDOptimizer
MDOptimizer(base: GradientTransformation, *, gain_lr: float = 0.001, gain_axis: Any = ('row', 'col'), gain_map: str = 'softplus')
Sentinel stored on a model via net.optimizer(jno.optimizers.md(...)).
Not an optax GradientTransformation — jno.core.solve() detects this object, builds the
real :func:md_decouple transform from it via :meth:build, and forces the model's outer
learning-rate scale to 1.0 so the nonlinear W_new − W update is applied intact. The
direction learning rate η_W is carried by the base optimizer you pass.
See :func:md for the public constructor and the list of not-yet-implemented features above
this class.
MMAOptimizer
MMAOptimizer(move: float = 0.2, lower: Optional[float] = None, upper: Optional[float] = None, asy_init: float = 0.5, asy_shrink: float = 0.7, asy_grow: float = 1.2, raa: float = 1e-05, dual_iters: int = 400, move_gamma: float = 1.0, move_min: float = 0.0)
Sentinel stored on a model by p.optimizer(jno.optimizers.mma(...)).
Not an optax GradientTransformation. jno.core.solve() detects it, injects
optax.sgd(1.0) as the transform that applies the step, and prepends a single
:class:MMACallback covering every model marked this way -- MMA solves one joint
subproblem, because the design variables are coupled through the shared constraints.
See :func:mma for the public constructor.
geometric_decay
geometric_decay(param: Any, gamma: float, *, start: float = 1.0, minimum: float = 0.0) -> GeometricDecay
Decay a scalar parameter geometrically, b_iter = gamma * b_(iter-1) — eq. (41).
Written for the perimeter barrier's weight (Jung, Yun & Kim, Computers & Structures 331
(2026) 108403, eq. 40-41). The objective is C - beta * P* log(P* - P): with beta fixed,
the barrier holds the design well short of the target P* forever, because the penalty is
still paying for perimeter it is allowed to spend. Shrinking beta lets the design approach
P* from below as the run proceeds — the paper's own words, "to reduce the approximation
error of the perimeter constraint", and the standard interior-point continuation.
Give the parameter optax.sgd(1.0) and initialise it to start; this hook writes it
directly, the same way :class:SIMPContinuation writes penal.
| PARAMETER | DESCRIPTION |
|---|---|
param
|
The scalar
TYPE:
|
gamma
|
Decay factor per iteration, in
TYPE:
|
start
|
Initial value. Must match the parameter's initializer.
TYPE:
|
minimum
|
Floor, so the barrier never vanishes entirely and let the design cross
TYPE:
|
heaviside_continuation
heaviside_continuation(param: Any, *, start: float = 1.0, maximum: float = 16.0, over: int, hold: int = 0, schedule: str = 'geometric') -> HeavisideContinuation
Sharpen a density projection over a declared number of iterations.
A smoothed-Heaviside projection (see :func:jno.np.heaviside) is only useful if its sharpness
rises: at beta = 1 it barely projects, and starting at the final beta gives the optimiser
a near-step map to descend through and it stalls in the first grey design it finds. The standard
answer is continuation, and this is it.
The ramp is keyed to this callback's own invocation count, never to epoch. That is the
whole design and it is not a detail: epoch is per-solve()-call, so a schedule reading it
ramps once per CHUNK and the physics then depends on how the driver loop was written. Measured on
a 3-D bracket, the same 250-iteration run reached beta 10.8 at CHUNK=10 and 1.6 at
CHUNK=50, with M_nd 0.008 against 0.078 -- two different optimisations wearing the same
configuration. Counting invocations is chunk-independent by construction: solve(1) five times
and solve(5) once give the identical ramp.
over is required and cannot be inferred. on_solve_begin does not carry the run
length, and a driver that also stops on a wall clock does not know it either until it stops --
so any inference would be wrong exactly where it matters, on the short run. Pass the
fraction-of-the-run form directly, over=int(0.8 * ITERS), which is what the schedule means.
Register it BEFORE any callback that reads the projected density, so they see the current value:
callbacks fire in list order, each seeing the previous one's grads. With
:func:simp_continuation that is::
beta = jno.np.parameter((1,), name="beta")
beta.dtype(jnp.float64)
beta.initialize(lambda k, sh, dtype=None: jnp.full(sh, 1.0))
beta.optimizer(optax.sgd(1.0))
hv = jno.optimizers.heaviside_continuation(beta, maximum=16.0, over=int(0.8 * ITERS))
sc = jno.optimizers.simp_continuation(penal, rho, physical=lambda r: hv.project(patch(r)))
crux.solve(CHUNK, callbacks=[hv, sc])
Give the parameter optax.sgd(1.0) and initialise it to start; this hook writes it
directly, the same way :class:SIMPContinuation writes penal.
Scope note. penal and beta both binarise, by different means -- SIMP makes grey
expensive, projection makes it unreachable. The paper's schedule raises penal only. Running
both is legitimate and is what a Heaviside topology optimisation usually does, but they compound,
so watch M_nd rather than assuming the two are independent.
| PARAMETER | DESCRIPTION |
|---|---|
param
|
The scalar
TYPE:
|
start
|
Initial sharpness. Must match the parameter's initializer.
TYPE:
|
maximum
|
Final sharpness, held once the ramp saturates.
TYPE:
|
over
|
Iterations from
TYPE:
|
hold
|
Iterations at
TYPE:
|
schedule
|
TYPE:
|
| ATTRIBUTE | DESCRIPTION |
|---|---|
value |
The current sharpness.
|
progress |
Position along the ramp,
|
saturated |
Whether
|
history |
One entry per invocation.
|
simp_continuation
simp_continuation(penal: Any, density: Any, *, start: float = 3.0, step: float = 1.0, maximum: float = 8.0, tol: float = 0.0001, window: int = 3, mnd_tol: float = 0.0001, physical: Optional[Callable] = None, watch: Optional[int] = None, every: int = 1, patience: Optional[int] = None) -> SIMPContinuation
Raise the SIMP exponent once the objective settles and the design is still grey.
Register it after the design optimiser's own callback -- jno.core applies
on_before_update hooks in list order, each seeing the previous one's gradients, and this
one only replaces the penal entry.
| PARAMETER | DESCRIPTION |
|---|---|
penal
|
The
TYPE:
|
density
|
The design-density parameter, read to compute
TYPE:
|
start
|
Initial exponent. 3 is the paper's, and the field's, default.
TYPE:
|
step
|
Increment per continuation step.
TYPE:
|
maximum
|
Ceiling, so a design that never binarises cannot run the exponent away.
TYPE:
|
tol
|
Relative-change tolerance for "the objective has converged" (the paper's 1e-4).
TYPE:
|
window
|
How many consecutive intervals must all be below
TYPE:
|
mnd_tol
|
Grey-level target. Above this, the exponent keeps rising. 1e-4, which is what the paper states in Sec. 3.1 -- 'if the gray level indicator M_nd remains larger than 10^-4, the penalization factor is increased by one until the criterion on M_nd is satisfied'. This default was 1e-2, attributed to the paper in the same breath; a design at M_nd = 1e-2 still carries visible intermediate-density members, which is the difference between their Fig. 4d and Fig. 4e.
TYPE:
|
physical
|
Optional map from the design density to the physical one, e.g.
TYPE:
|
watch
|
Index into
TYPE:
|
every
|
Sample the convergence window every
TYPE:
|
patience
|
Raise
TYPE:
|
| ATTRIBUTE | DESCRIPTION |
|---|---|
history |
One
|
m_nd |
The grey-level indicator (eq. 20) at the last sample, on the physical density.
|
Example::
penal = jno.np.parameter((1,), name="penal")
penal.initialize(lambda k, s, dtype=None: jnp.full(s, 3.0))
penal.optimizer(optax.sgd(1.0))
cont = jno.optimizers.simp_continuation(penal, rho)
jno.core([...], domain=d).solve(300, callbacks=[cont])
for epoch, penal_new, why in cont.history:
print(epoch, penal_new, why) # when each step fired, and what fired it
References
Jung, Yun & Kim, Computers & Structures 331 (2026) 108403, Sec. 2.3.2, Fig. 4a. Sigmund, Struct. Multidisc. Optim. 21(2), 2001, 120-127 (SIMP and continuation).
md
md(base: GradientTransformation, *, gain_lr: float = 0.001, gain_axis: Any = ('row', 'col'), gain_map: str = 'softplus') -> MDOptimizer
Return a Magnitude–Direction Decoupling sentinel for net.optimizer().
Magnitude–Direction (MD) Decoupling (Hägele, Hernández-Cano, Kosson & Jaggi, 2026,
arXiv:2606.25971, Algorithm 2) factorizes each weight matrix W = diag(γ_row) Ŵ diag(γ_col)
into a fixed-Frobenius-norm direction Ŵ on a hypersphere and learnable per-row/per-column
magnitude gains, updated at separate rates. It removes the need for weight decay and warmup and
transfers the optimal learning rate across model width. The split lives entirely inside the
optimizer — the model still sees one fused weight tensor.
Unlike the bare :func:md_decouple transform, this sentinel form lets jno.core force the
model's outer learning-rate scale to 1.0 automatically, so the nonlinear W_new − W
update is never corrupted — you do not need a separate net.scale(1.0) call. The
direction learning rate η_W is the LR of the base optimizer you pass.
| PARAMETER | DESCRIPTION |
|---|---|
base
|
the direction optimizer — any pure-gradient optax
TYPE:
|
gain_lr
|
learning rate η_γ of the separate Adam that steps the magnitude gains.
TYPE:
|
gain_axis
|
TYPE:
|
gain_map
|
positive map for the gains; only
TYPE:
|
Example::
import optax, jno
net.optimizer(jno.optimizers.md(optax.adam(1e-3), gain_lr=1e-3))
crux = jno.core(losses)
crux.solve(2000)
References
Hägele, Hernández-Cano, Kosson & Jaggi, Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors, 2026, arXiv:2606.25971, Algorithm 2.
scale_by_soap
scale_by_soap(b1: float = 0.95, b2: float = 0.95, shampoo_beta: float = -1, eps: float = 1e-08, correct_bias: bool = True, precondition_frequency: int = 10, max_precond_dim: int = 10000, precondition_1d: bool = False, precision: PrecisionLike = jax.lax.Precision.HIGHEST, mu_dtype: Optional[Any] = None, qr_dtype: Any = jnp.float32) -> GradientTransformation
Implements SOAP algorithm (https://arxiv.org/abs/2409.11321). Based on the original implementation at https://github.com/nikhilvyas/SOAP.
| PARAMETER | DESCRIPTION |
|---|---|
b1
|
Adam's beta1 parameter. Defaults to 0.95.
TYPE:
|
b2
|
Adam's beta2 parameter. Defaults to 0.95.
TYPE:
|
shampoo_beta
|
If >= 0, use this beta for the preconditioner (
TYPE:
|
eps
|
Adam's epsilon for numerical stability. Defaults to 1e-8.
TYPE:
|
correct_bias
|
Whether to use bias correction for the Adam moments. Defaults to True.
TYPE:
|
precondition_frequency
|
How often to update the preconditioner. Defaults to 10.
TYPE:
|
max_precond_dim
|
Maximum dimension of the preconditioner. Set to 10000 to exclude most common vocab sizes while including layers. Defaults to 10000.
TYPE:
|
precondition_1d
|
Whether to precondition 1D gradients. If False, 1D params use Adam-style updates. Defaults to False.
TYPE:
|
precision
|
Precision to use. Defaults to jax.lax.Precision.HIGHEST.
TYPE:
|
mu_dtype
|
dtype for the first and second moment estimates (exp_avg and exp_avg_sq). If None, uses the same dtype as the parameters. Useful for mixed-precision training. Defaults to None.
TYPE:
|
qr_dtype
|
dtype used for eigen/QR computations and preconditioner storage. Defaults to float32 to avoid float64 upcasts in mixed precision.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
GradientTransformation
|
The SOAP gradient transformation.
TYPE:
|
scale_by_ss_quasi_newton
scale_by_ss_quasi_newton(method: str = 'ssbfgs', scale_init_precond: bool = True) -> base.GradientTransformation
Scale updates by the Self-Scaled Broyden/BFGS inverse Hessian.
This maintains a dense approximation of the inverse Hessian
:math:H_k and returns the preconditioned gradient
:math:H_k \nabla f(w_k). Unlike L-BFGS, the full
:math:n \times n matrix is stored, so this is suitable only for
small to medium scale problems.
The inverse Hessian is updated using a self-scaled formula:
.. math::
H_{k+1} = \frac{1}{\tau_k}(H_k - \frac{H_k y_k y_k^\top H_k}{y_k^\top H_k y_k} + \phi_k v_k v_k^\top) + \frac{s_k s_k^\top}{y_k^\top s_k}
where the self-scaling factors :math:\tau_k and :math:\phi_k depend on
the chosen method ('ssbfgs' or 'ssbroyden').
This function is typically chained with a line search such as
:func:optax.scale_by_zoom_linesearch.
| PARAMETER | DESCRIPTION |
|---|---|
method
|
TYPE:
|
scale_init_precond
|
if
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
A
|
class:
TYPE:
|
References
Urbán et al, Unveiling the optimization process of physics
informed neural networks: How accurate and competitive can
PINNs be?
<https://doi.org/10.1016/j.jcp.2024.113656>_, 2025
.. warning::
This optimizer stores a dense :math:n \\times n matrix
where :math:n is the total number of parameters. It is
memory intensive and best suited for small to medium scale
problems.
ssbfgs
ssbfgs(learning_rate: Optional[ScalarOrSchedule] = None, scale_init_precond: bool = True, linesearch: Optional[Union[GradientTransformationExtraArgs, GradientTransformation]] = _linesearch.scale_by_zoom_linesearch(max_linesearch_steps=20, initial_guess_strategy='one')) -> base.GradientTransformationExtraArgs
Self-Scaled BFGS optimizer.
SSBFGS is a quasi-Newton method that maintains a dense approximation of
the inverse Hessian. Unlike L-BFGS which uses a limited-memory
approximation, this method stores the full :math:n \times n matrix and
updates it using the self-scaled BFGS formula, yielding improved scaling
of the search direction at the cost of higher memory usage.
The inverse Hessian :math:H_k is updated each iteration as:
.. math::
H_{k+1} = \frac{1}{\tau_k}!\left(H_k - \frac{H_k y_k y_k^\top H_k}{y_k^\top H_k y_k} + v_k v_k^\top\right) + \frac{s_k s_k^\top}{y_k^\top s_k}
where the self-scaling factor :math:\tau_k uses the BFGS formula
(:math:\phi_k = 1).
| PARAMETER | DESCRIPTION |
|---|---|
learning_rate
|
optional global scaling factor, either fixed or evolving along iterations with a scheduler. By default the learning rate is handled by the line search.
TYPE:
|
scale_init_precond
|
whether to scale the initial identity preconditioner by a capped reciprocal of the gradient norm.
TYPE:
|
linesearch
|
an instance of
:class:
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
A
|
class:
TYPE:
|
Examples:
>>> import optax
>>> import jax
>>> import jax.numpy as jnp
>>> def f(x): return jnp.sum(x ** 2)
>>> solver = optax.contrib.ssbfgs()
>>> params = jnp.array([1., 2., 3.])
>>> print('Objective function: ', f(params))
Objective function: 14.0
>>> opt_state = solver.init(params)
>>> value_and_grad = optax.value_and_grad_from_state(f)
>>> for _ in range(5):
... value, grad = value_and_grad(params, state=opt_state)
... updates, opt_state = solver.update(
... grad, opt_state, params,
... value=value, grad=grad, value_fn=f,
... )
... params = optax.apply_updates(params, updates)
References
Urbán et al, Unveiling the optimization process of
physics informed neural networks: How accurate and
competitive can PINNs be?
<https://doi.org/10.1016/j.jcp.2024.113656>_, 2025
.. warning::
This optimizer stores a dense :math:n \\times n matrix
where :math:n is the total number of parameters. It is
memory intensive and best suited for small to medium
scale problems.
.. warning::
This optimizer works best with a line search (default is a zoom line
search). See the example above for best use in a non-stochastic setting
where gradients computed by the line search can be recycled using
:func:optax.value_and_grad_from_state.
.. seealso:: :func:optax.contrib.ssbroyden
Bayesian inference (jno.bayesian)
Backs model.bayesian(...) / model.vi(...); see Bayesian Sampling.
jno.bayesian
Adapters that wire blackjax MCMC / VI kernels into jno's training loop.
Most of this module is internal — users compose blackjax kernels
directly and attach them per parameter via :meth:jno.trace.Model.bayesian
or :meth:jno.trace.Model.vi. A small public surface lives at
jno.bayesian.*: convergence diagnostics (:func:rhat, :func:ess)
and the logdensity-aware initializer protocol — currently
:class:PathfinderInitializer (:func:pathfinder),
:class:LaplaceInitializer (:func:laplace), and
:class:SVGDInitializer (:func:svgd). The training loop in
:mod:jno.core dispatches each model's per-step update either through
:mod:optax or through the blackjax kernel configured here.
Supported kernel families (duck-typed via the first argument name of the
factory's .differentiable callable):
-
Full-data kernels — first argument
logdensity_fn. Covers :func:blackjax.nuts(Hoffman & Gelman 2014), :func:blackjax.hmc, :func:blackjax.mala, and any other kernel that takes a log-density on the full dataset. Step signature:kernel.step(rng_key, state)→(new_state, info). -
SG-MCMC kernels — first argument
grad_estimator. Covers :func:blackjax.sgld(Welling & Teh 2011) and :func:blackjax.sghmc. Step signature:kernel.step(rng_key, state, minibatch, step_size)→new_state.
References
Hoffman, M. D., & Gelman, A. (2014). The No-U-Turn Sampler. Journal of Machine Learning Research, 15(1), 1593-1623. Welling, M., & Teh, Y. W. (2011). Bayesian Learning via Stochastic Gradient Langevin Dynamics. In ICML 2011, pp. 681-688.
LaplaceInitializer
dataclass
LaplaceInitializer(map_steps: int = 500, map_optimizer: Any = None, hessian_strategy: str = 'diagonal', ridge: float = 1e-06)
Laplace approximation to the posterior — warm-start a chain from the MAP plus its local Gaussian approximation.
Three-step algorithm:
- Find the maximum a-posteriori (MAP) point by optimising
-log p(theta | data)with an optax optimiser (Adam by default). The optimisation runs as a JIT-compiledjax.lax.scanovermap_stepsiterations. -
Compute the Hessian
H = -∇²log pat the MAP. Two strategies: -
hessian_strategy="full"— full(D, D)Hessian via :func:jax.hessian. Numerically clean but memory cost grows asD². Right choice for small models / scalar PDE coefficients (D up to ~1000). -
hessian_strategy="diagonal"(default) — only the diagonal of the Hessian is computed via D Hessian-vector probes. Memory cost isO(D)instead ofO(D²)— required for BNN-scale problems. Compute cost is comparable to "full" but peak memory is much smaller because no D×D matrix is ever materialised. -
Approximate the posterior as
N(MAP, H⁻¹). Fornum_chains=1the warm position is the MAP; fornum_chains>1we draw K i.i.d. samples from this Gaussian (proper over-dispersion). The diagonal ofH⁻¹is returned as the kernel'sinverse_mass_matrix.
A small ridge (default 1e-6) is added to the diagonal of H
before any inversion / Cholesky to guard against rank-deficient
Hessians at non-converged MAP estimates.
Trade-offs vs :class:PathfinderInitializer
- Pathfinder uses L-BFGS — explores the optimisation path and produces a normal approximation from the path's inverse-Hessian factors. Often robust on multi-modal or curved posteriors; cheaper than computing a Hessian for very large models.
- Laplace uses gradient descent on
-log pand the exact Hessian at the MAP. More accurate locally if the posterior is well-approximated by a Gaussian; falls back to ridge-regularised sampling if H is ill-conditioned.
References
MacKay, D. J. C. (1992). A Practical Bayesian Framework for Backpropagation Networks. §6 (Laplace approximation around the posterior mode). Neural Computation, 4(3), 448-472. https://doi.org/10.1162/neco.1992.4.3.448
Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., & Hennig, P. (2021). Laplace Redux — Effortless Bayesian Deep Learning. §2 (Laplace approximations for neural networks). Advances in Neural Information Processing Systems (NeurIPS). https://arxiv.org/abs/2106.14806
Magnani, E., Krämer, N., Pförtner, M., & Hennig, P. (2024). Linearization Turns Neural Operators into Function-Valued Gaussian Processes. §3 (linearised-Laplace for neural operators). https://arxiv.org/abs/2406.05072
PathfinderInitializer
dataclass
PathfinderInitializer(maxiter: int = 30, num_samples: int = 200, maxcor: int = 10, imm_estimator_samples: int = 500, lbfgs_kwargs: dict = dict())
Warm-start a chain via :func:blackjax.pathfinder.
Runs L-BFGS on the log-density; the inverse-Hessian factors along
the optimisation path are turned into a normal approximation to the
posterior. From that fitted q we draw (a) the warm starting
position (the MAP-ish state.position for K=1; K distinct
samples for K>1 — proper over-dispersion at zero extra cost)
and (b) per-dimension variance estimates that become the kernel's
diagonal inverse_mass_matrix.
Reference
Zhang, L., Carpenter, B., Gelman, A., & Vehtari, A. (2022). Pathfinder: Parallel quasi-Newton variational inference. Journal of Machine Learning Research, 23(306), 1-49. https://arxiv.org/abs/2108.03782
SVGDInitializer
dataclass
SVGDInitializer(num_iters: int = 500, num_particles: int | None = None, optimizer: Any = None, kernel: Callable | None = None, init_jitter: float | None = None)
Stein Variational Gradient Descent warm-start.
Runs SVGD (Liu & Wang 2016) — a deterministic, particle-based variational inference algorithm — and uses the final particle cloud as the warm-start. Each particle is dragged toward the posterior by a kernelised functional gradient that combines a repulsive RBF term (keeps particles diverse) with an attractive log-density gradient term (pulls each particle toward higher posterior density). At convergence, the particles approximate the posterior distribution.
Three-step algorithm:
- Initialise
num_particlesparticles by perturbing the user-supplied position with Gaussian noise of stdinit_jitter. Wheninit_jitteris left at its default (None), jno picksmax(0.1 * std(position), 1e-3)so particles start one-tenth of a parameter-scale apart — a sensible "small but visible" spread that adapts to whatever initialisation the model uses. Explicit positive floats are respected as absolute std values. Defaultnum_particles = max(num_chains, 32)so we always have at least 32 particles for the variance estimate, even when the caller only asked for 1 chain. - Run
num_itersSVGD steps using :func:blackjax.svgdwith the supplied optax optimiser (Adam by default) and the default RBF kernel (overridable viakernel). The whole run is wrapped in :func:jax.lax.scanfor fast XLA-side iteration. -
Return:
-
For
num_chains == 1— the particle-cloud mean as the warm starting position (most stable summary). - For
num_chains > 1— the firstnum_chainsparticles as K distinct warm starting positions. The particle dynamics already provide proper over-dispersion; no extra jitter step is needed.
The per-dimension variance of the particle cloud is
returned as the diagonal inverse_mass_matrix (a cheap
empirical-Bayes estimate of the posterior covariance).
Trade-offs vs Pathfinder / Laplace
- Pathfinder approximates the posterior via L-BFGS factors — cheap and unimodal.
- Laplace uses the exact Hessian at the MAP — accurate locally but ignores posterior modes beyond the MAP basin.
- SVGD captures multi-modal structure when present: the
repulsive RBF kernel pushes particles apart so multiple modes
can be discovered with enough particles. Cost grows with
num_particles²per step (pairwise kernel interactions).
Reference
Liu, Q., & Wang, D. (2016). Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm. §3 (the SVGD update rule). Advances in Neural Information Processing Systems (NeurIPS), 29, 2378-2386. https://arxiv.org/abs/1608.04471
rhat
Rank-normalised, folded R-hat (Gelman & Rubin 1992, improved by Vehtari et al. 2021).
chain is shaped (K, N, *param) — K chains, N draws each.
Returns an array of shape *param with one R-hat per parameter
component. R-hat ≈ 1.0 indicates the chains are exploring the same
distribution; values above ~1.05 suggest non-convergence.
Parameters
chain : array
(K, N, *param) posterior chain.
strategy : {"auto", "multichain", "split"}, default "auto"
How to estimate R-hat:
* ``"auto"`` (default): split-R-hat when K == 1, multichain
R-hat when K >= 2 — preserves the historical behaviour.
* ``"multichain"``: require K >= 2; raise ``ValueError`` if not.
Use when you've explicitly run multiple chains and want a
loud failure if the layout doesn't carry them.
* ``"split"``: split every chain in half and treat the halves
as 2K independent chains. Useful as a stationarity check
on top of cross-chain mixing for K >= 2.
Notes
For K=1 the default "auto" strategy falls back to single-chain
split-R-hat (still informative for stationarity but not as strong
as a true multi-chain diagnostic). Pass strategy="multichain"
to forbid this silent fallback.
References
Gelman, A., & Rubin, D. B. (1992). Inference from iterative simulation using multiple sequences. Statistical Science 7(4), 457-511.
Vehtari, A., Gelman, A., Simpson, D., Carpenter, B., & Bürkner, P.-C. (2021). Rank-Normalization, Folding, and Localization: An Improved R̂ for Assessing Convergence of MCMC. Bayesian Analysis 16(2), 667-718.
ess
Effective sample size via FFT-based autocorrelation, averaged across chains.
chain is shaped (K, N, *param). Returns an array of shape
*param with one ESS per parameter component. Values are bounded
above by K * N (total raw samples). ESS larger than ~100 per
parameter is typically considered sufficient.
Implementation: per-chain autocovariance via real FFT, averaged across chains, summed with the initial-monotone-sequence truncation rule of Geyer (1992).
References
Gelman, A., Carlin, J. B., Stern, H. S., Dunson, D. B., Vehtari, A., & Rubin, D. B. (2014). Bayesian Data Analysis, 3rd ed., Chapman & Hall/CRC, §11.5.
default_gaussian_prior
Wide isotropic Gaussian log-prior over a pytree position.
log p(θ) = -‖θ‖² / (2σ²), summed over every inexact-array leaf.
This is the internal fallback when .bayesian(prior=None). Users
composing custom priors should prefer the named factories in
:data:priors (e.g. jno.bayesian.priors.gaussian(sigma=5.0))
— they share the same logdensity-over-pytree contract but document
their arguments and validate inputs at construction time.
laplace
Build a :class:LaplaceInitializer for use with
:meth:jno.trace.Model.initialize.
Usage::
net.initialize(jno.bayesian.laplace(map_steps=500,
hessian_strategy="diagonal"))
net.bayesian(blackjax.nuts, step_size=1e-2, warmup=100, adapt=True)
Kwargs (see :class:LaplaceInitializer for full descriptions):
map_steps(default 500) — number of optimiser iterations to find the MAP.map_optimizer(defaultoptax.adam(1e-2)) — any optaxGradientTransformation.hessian_strategy(default"diagonal") —"diagonal"computes diag(H) via D HVPs (O(D) memory),"full"materialises the (D, D) Hessian.ridge(default1e-6) — diagonal stabiliser added to H before inversion.
Reference
MacKay 1992 §6; Daxberger et al. 2021 §2. See
:class:LaplaceInitializer for full citations.
pathfinder
Build a :class:PathfinderInitializer for use with
:meth:jno.trace.Model.initialize.
Usage::
net.initialize(jno.bayesian.pathfinder(maxiter=30, num_samples=200))
net.bayesian(blackjax.nuts, step_size=1e-2, warmup=100, adapt=True)
All kwargs are forwarded to the underlying
blackjax.pathfinder.approximate call (maxiter, num_samples,
maxcor, plus any lbfgs_kwargs-eligible LBFGS knobs).
Noise nodes (jno.noise)
jno.noise._NoiseNamespace
Namespace exposed as jno.noise.
gaussian
Zero-mean Gaussian noise ~ N(0, std²).
Parameters
std : float
Standard deviation of the noise.
ndim : int
Number of output dimensions (last axis size). Use ndim=2 for
2-D vector noise, ndim=3 for 3-D, etc. Default 1.
uniform
Uniform noise ~ U(low, high).
Parameters
low : float
Lower bound of the uniform distribution.
high : float
Upper bound of the uniform distribution.
ndim : int
Number of output dimensions. Default 1.
laplace
Zero-mean Laplace noise with the given standard deviation.
Parameters
std : float
Standard deviation (scale = std / sqrt(2)).
ndim : int
Number of output dimensions. Default 1.
grf
grf(*coords: Any, length_scale: float = 0.1, variance: float = 1.0, kernel: str = 'matern', nu: float = 1.5, modes: int = 256, ndim: int = 1) -> FunctionCall
A Gaussian random field over coords, redrawn every training step.
Unlike the pointwise distributions above, a GRF is spatially correlated: nearby points
get nearby values, with the correlation set by length_scale. That makes it an input
function, which is what turns operator learning from a parametric fit into an operator --
a fresh in-distribution input every step, with no dataset and no data API::
f = jno.noise.grf(x, y, length_scale=0.1)
u = net(f).scalar.bind(x=x, y=y)
crux = jno.core([(u.laplacian(x, y) + f).mse])
Algorithm -- spectral representation (a.k.a. the randomization method / random Fourier features):
.. math:: f(x) = \sqrt{\sigma^2/M} \sum_{j=1}^{M} \left[ a_j \cos(\omega_j \cdot x) + b_j \sin(\omega_j \cdot x) \right], \qquad a_j, b_j \sim \mathcal{N}(0, 1),
with the frequencies :math:\omega_j drawn from the kernel's spectral density. Conditional
on :math:\{\omega_j\} this is exactly a zero-mean Gaussian field, and its covariance
converges to the kernel at :math:O(M^{-1/2}).
For Matern-:math:\nu the spectral density is a multivariate Student-t with :math:2\nu
degrees of freedom and scale :math:I/\ell^2, sampled as its Gaussian scale mixture
:math:\omega = z / (\ell\sqrt{g}) with :math:z \sim \mathcal{N}(0, I) and
:math:g \sim \mathrm{Gamma}(\nu, 1/\nu). The squared-exponential kernel is the
:math:\nu \to \infty limit, where :math:g \to 1.
References
Shinozuka & Deodatis, Simulation of Stochastic Processes by Spectral Representation, Appl. Mech. Rev. 44(4) (1991), sections 2-3 -- the spectral representation. Rahimi & Recht, Random Features for Large-Scale Kernel Machines, NIPS 20 (2007), section 2 -- the Monte-Carlo frequency draw. Rasmussen & Williams, Gaussian Processes for Machine Learning (2006), section 4.2 eq. 4.15 -- the Matern spectral density.
| PARAMETER | DESCRIPTION |
|---|---|
*coords
|
The coordinate Variables the field varies over, e.g.
TYPE:
|
length_scale
|
Correlation length. Much larger than the domain gives a near-constant field; much smaller than the point spacing gives near-white noise.
TYPE:
|
variance
|
Marginal variance :math:
TYPE:
|
kernel
|
TYPE:
|
nu
|
Matern smoothness.
TYPE:
|
modes
|
Number of spectral modes :math:
TYPE:
|
ndim
|
Number of independent field components (the trailing axis).
TYPE:
|
Limitations
An approximate GP sample: exact only in the limit :math:M \to \infty, with
covariance error :math:O(M^{-1/2}). Exact circulant embedding (Dietrich & Newsam,
SIAM J. Sci. Comput. 18(4), 1997) is not implemented because it needs a regular
grid, whereas the evaluator sees a flat point cloud. Components share one frequency
draw, so they share its finite-M error while remaining independent fields.
Like every noise node the realisation is fixed across the timesteps of one window.
Units & non-dimensionalization (jno.units)
Annotate the dimension and characteristic magnitude of a leaf with .unit(...) / .scale(...), then
audit consistency, extract the dimensionless groups (Fourier, Péclet, …) of a residual, and rewrite
it to a well-scaled O(1) form. Worked usage: Operations → Units &
non-dimensionalization.
jno.trace.units
Lightweight physical-unit algebra for graph-time dimensional analysis.
Units are graph-level metadata: they live on jno's Python wrapper objects
(:class:jno.trace.Placeholder) and are inferred at trace time, before any
JAX computation, so a user can audit that a PDE formulation is dimensionally
consistent. See :mod:jno.trace.unit_log for the inference walk and the
public jno.units.check(...) entry point.
pint is intentionally not a dependency — jno minimises external deps and
the algebra needed here is small. A :class:Unit is just a map from SI base
dimension to a :class:fractions.Fraction exponent (fractions so sqrt
stays exact). The seven SI bases plus a handful of common derived aliases are
recognised by :meth:Unit.parse.
Unit
A physical unit as a map of SI base dimension → rational exponent.
Instances are immutable value objects: equality and hashing are by the
(zero-stripped) exponent map, so Unit.parse("m/s") == Unit.parse("m s^-1").
parse
classmethod
Parse a unit string such as 'm', 'm/s', 'kg/m^3', 'Pa'.
Grammar (informal): a numerator optionally followed by / and a
denominator; each is a juxtaposition of base[^power] factors
separated by spaces, *, or ·. Powers may be negative or
fractional (m^1/2). Recognises the 7 SI bases plus the aliases in
:data:_ALIASES (T→K, Pa, N, J, W, Hz, …).
Dimensionless is spelled '', '1', or '-'.
Rescaler
Back-map between physical and dimensionless (O(1)) coordinates/fields.
Produced by :func:rescale. Holds the characteristic length(s) L of
each coordinate and the field scale U so a rescaled problem can be set
up and its solution recovered: u_physical(x) = U · û(x / L).
rescaled_context
Return a new context dict with each coordinate column scaled to O(1).
rescaled_domain
Return a shallow copy of domain whose context is rescaled to O(1).
The user's domain is never mutated.
NondimReport
Result of :func:nondimensionalize: one :class:ResidualReport each.
UnitLogger
check
Infer and audit physical units across one or more expression trees.
Walks each node in nodes (a single :class:~jno.trace.Placeholder or a
list/tuple of them) in post-order, inferring a unit for every sub-node and
recording any dimensional inconsistencies. Optionally writes a
human-readable log.
Parameters
nodes:
A traced expression or a sequence of them (e.g. the residuals passed to
jno.fem([...])).
log:
If given, a filesystem path the per-node log is written to.
Returns
UnitLogger
Inspect .entries ((label, unit, warnings) triples) or
.warnings (flat list) programmatically.
infer
Convenience: return the inferred unit of a single expression's root.
Equivalent to check([node]).entries[-1][1] — runs a fresh, non-mutating
walk and hands back the top node's unit (None if not derivable).
nondimensionalize
Derive characteristic magnitudes and dimensionless groups for residual(s).
For each residual in nodes, splits the top-level additive terms, infers
each term's :class:~jno.trace.units.Unit (via :func:infer) and
characteristic magnitude (via the per-leaf .scale walk), then forms the
dimensionless group of each term as πᵢ = Sᵢ / S_ref.
Phase A only — this is pure analysis and does not change what the solver computes.
Parameters
nodes: A residual expression or a sequence of them. ref: Index of the reference term whose magnitude normalises the groups. Defaults to the first term with both a known unit and a known scale. report: If given, a path the human-readable report is written to.
Returns
NondimReport
.residuals[i].terms[j].pi holds each term's dimensionless group.
rescale
Transform residual constraint(s) into a solvable dimensionless form.
Multiplies each additive term of every residual by gᵢ = Kᵢ / S_ref so
that, evaluated on the :class:Rescaler's O(1) coordinate context (and with
the network/trial output representing the dimensionless field), the residual
becomes Σ signᵢ·πᵢ·term̂ᵢ — O(1) and dimensionless, with the
dimensionless groups (Fourier/Péclet) as the leading coefficients.
Phase B — opt-in. Returns (transformed_constraints, rescaler); nothing
is solved and the input domain is never mutated.
Parameters
constraints:
A residual expression or a sequence of them.
domain:
Unused for the transform itself; accepted so callers can pair it with
rescaler.rescaled_domain(domain).
Returns
(list, Rescaler)
The transformed constraints (same length/order as the input) and the
back-map :class:Rescaler.
Adaptive resampling (jno.sampler)
Residual-adaptive collocation strategies — see Adaptive Resampling.
jno.utils.adaptive.resampling
Base class for resampling strategies.
sampler
Factory class for creating resampling strategies.
random
staticmethod
Random resampling - simple baseline to prevent overfitting.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Fraction of points to resample
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
rad
staticmethod
rad(resample_every: int = 100, resample_fraction: float = 0.1, start_epoch: int = 1000, k: int = 10)
Residual-based Adaptive Distribution (RAD) resampling.
Focuses on regions with high PDE residuals (high errors).
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Fraction of points to resample
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
k
|
Number of top residual points to cluster around
TYPE:
|
rard
staticmethod
rard(resample_every: int = 100, resample_fraction: float = 0.1, start_epoch: int = 1000, power: float = 2.0)
Residual-based Adaptive Refinement with Distribution (RARD).
Uses importance sampling based on residual^power distribution.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Fraction of points to resample
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
power
|
Power for residual-based importance weighting
TYPE:
|
ha
staticmethod
ha(resample_every: int = 100, resample_fraction: float = 0.5, start_epoch: int = 1000, alternate: bool = True, random_first: bool = True)
Hybrid Adaptive (HA) resampling strategy.
Alternates between random and adaptive phases for regularization while focusing on high-error regions.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Fraction of points to retain in adaptive phase
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
alternate
|
Whether to alternate between random and adaptive phases
TYPE:
|
random_first
|
Start with random phase if True
TYPE:
|
r3
staticmethod
r3(resample_every: int = 100, resample_fraction: float = 0.7, start_epoch: int = 1000, threshold_mode: str = 'mean', min_keep_frac: float = 0.3, max_keep_frac: float = 0.9)
Residual-based Refinement and Resampling (R3).
Keeps points with high residuals and replaces low-residual points with new samples from the candidate pool.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Not used for R3 (uses threshold instead)
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
threshold_mode
|
Threshold for keeping points ("mean", "median", or float)
TYPE:
|
min_keep_frac
|
Minimum fraction of points to keep
TYPE:
|
max_keep_frac
|
Maximum fraction of points to keep
TYPE:
|
cr3
staticmethod
cr3(resample_every: int = 100, resample_fraction: float = 0.5, start_epoch: int = 1000, t_index: int = -1, alpha: float = 5.0, gamma0: float = -0.5, eta_g: float = 0.001, epsilon: float = 20.0, delta_max: float = 0.1, min_keep_frac: float = 0.1, max_keep_frac: float = 0.9)
Causal Retain-Resample (CR3) with time gating.
Applies causal time gate for time-dependent PDEs where causality matters (e.g., wave equations, diffusion).
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Target fraction to keep
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
t_index
|
Index of time column (default -1 for last column)
TYPE:
|
alpha
|
Gate steepness (default 5.0)
TYPE:
|
gamma0
|
Initial gate position (default -0.5)
TYPE:
|
eta_g
|
Gate learning rate (default 1e-3)
TYPE:
|
epsilon
|
Gate update damping (default 20.0)
TYPE:
|
delta_max
|
Maximum gate step (default 0.1)
TYPE:
|
min_keep_frac
|
Minimum fraction to keep (default 0.1)
TYPE:
|
max_keep_frac
|
Maximum fraction to keep (default 0.9)
TYPE:
|
pinnfluence
staticmethod
pinnfluence(resample_every: int = 500, resample_fraction: float = 0.2, start_epoch: int = 2000, alpha: float = 1.0, c: float = 1.0, candidate_factor: float = 3.0)
PINNFluence: Influence function-based adaptive sampling (simplified).
Uses gradient-based scoring to identify points with high potential impact on reducing the loss.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs (use larger values, ~500+)
TYPE:
|
resample_fraction
|
Fraction of points to replace
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
alpha
|
Score exponent for sampling (default 1.0)
TYPE:
|
c
|
Additive smoothing constant (default 1.0)
TYPE:
|
candidate_factor
|
Pool size multiplier (default 3.0)
TYPE:
|
ResamplingStrategy
Base class for collocation point resampling strategies.
Resampling strategies adaptively update the training points during optimization to focus computation on regions with high errors or interesting dynamics.
Initialize resampling strategy.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Fraction of points to resample (0.0 to 1.0)
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
should_resample
Check if resampling should occur at this epoch.
| PARAMETER | DESCRIPTION |
|---|---|
epoch
|
Current training epoch
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
bool
|
True if resampling should happen |
resample
abstractmethod
resample(points: ndarray, residuals: ndarray, domain: Any, tag: str, epoch: int, rng_key: ndarray, candidates: Optional[ndarray] = None) -> jnp.ndarray
Compute new sample points.
| PARAMETER | DESCRIPTION |
|---|---|
points
|
Current points (N, D)
TYPE:
|
residuals
|
Residual values at current points (N,) or (B, N)
TYPE:
|
domain
|
domain object for sampling new candidates
TYPE:
|
tag
|
domain tag being resampled
TYPE:
|
epoch
|
Current training epoch
TYPE:
|
rng_key
|
JAX random key
TYPE:
|
candidates
|
Pre-drawn candidate points (N_pool, D) from
Under
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
New points (N, D) |
RAD
RAD(resample_every: int = 100, resample_fraction: float = 0.1, start_epoch: int = 1000, k: int = 10)
Residual-based Adaptive Distribution (RAD) resampling.
Resamples points based on residual magnitude - focuses on regions with high PDE residuals (high errors).
Reference: Lu et al. "DeepXDE: A Deep Learning Library for Solving Differential Equations" (2021)
Initialize RAD resampling.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Fraction of points to resample
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
k
|
Number of candidate points to sample per replacement
TYPE:
|
resample
resample(points: ndarray, residuals: ndarray, domain: Any, tag: str, epoch: int, rng_key: ndarray, candidates: Optional[ndarray] = None) -> jnp.ndarray
Resample based on residual magnitude.
Removes points with lowest residuals, adds new points near high-residual regions.
| PARAMETER | DESCRIPTION |
|---|---|
points
|
Current points (N, D)
TYPE:
|
residuals
|
Residual magnitudes (N,) or (B, N)
TYPE:
|
domain
|
domain object
TYPE:
|
tag
|
domain tag
TYPE:
|
epoch
|
Current epoch
TYPE:
|
rng_key
|
JAX random key
TYPE:
|
candidates
|
Pre-drawn candidate points from draw_candidates()
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
New points (N, D) |
RARD
RARD(resample_every: int = 100, resample_fraction: float = 0.1, start_epoch: int = 1000, power: float = 2.0)
Residual-based Adaptive Refinement with Distribution (RARD).
Similar to RAD but uses importance sampling based on residual distribution. Samples new points proportionally to residual^p where p is a power parameter.
Initialize RARD resampling.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Fraction of points to resample
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
power
|
Power for residual-based importance weighting
TYPE:
|
resample
resample(points: ndarray, residuals: ndarray, domain: Any, tag: str, epoch: int, rng_key: ndarray, candidates: Optional[ndarray] = None) -> jnp.ndarray
Resample using residual-weighted importance sampling.
| PARAMETER | DESCRIPTION |
|---|---|
points
|
Current points (N, D)
TYPE:
|
residuals
|
Residual magnitudes (N,) or (B, N)
TYPE:
|
domain
|
domain object
TYPE:
|
tag
|
domain tag
TYPE:
|
epoch
|
Current epoch
TYPE:
|
rng_key
|
JAX random key
TYPE:
|
candidates
|
Pre-drawn candidate points from draw_candidates()
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
New points (N, D) |
CR3
CR3(resample_every: int = 100, resample_fraction: float = 0.5, start_epoch: int = 1000, t_index: int = -1, alpha: float = 5.0, gamma0: float = -0.5, eta_g: float = 0.001, epsilon: float = 20.0, delta_max: float = 0.1, min_keep_frac: float = 0.1, max_keep_frac: float = 0.9)
Causal Retain-Resample (CR3) with time gating.
Applies a causal time gate g(t; γ) to modulate residuals, encouraging the network to learn from earlier times first. The gate parameter γ evolves during training to progressively include later times.
This is particularly effective for time-dependent PDEs where causality is important (e.g., wave equations, diffusion).
Reference
"Causal PINN: Respecting causality in physics-informed neural networks" https://proceedings.mlr.press/v202/daw23a.html
Algorithm
- Compute residual rI at interior points
- Apply time gate: g(t) = 0.5 * (1 - tanh(α * (t_norm - γ)))
- Score F = rI * g(t), keep points where F > mean(F)
- Update γ causally: γ += η_g * min(exp(-ε * Lg), δ_max)
Initialize CR3 resampling.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Target fraction to keep
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
t_index
|
Index of time column (default -1 for last column)
TYPE:
|
alpha
|
Gate steepness (default 5.0)
TYPE:
|
gamma0
|
Initial gate position (default -0.5)
TYPE:
|
eta_g
|
Gate learning rate (default 1e-3)
TYPE:
|
epsilon
|
Gate update damping (default 20.0)
TYPE:
|
delta_max
|
Maximum gate step (default 0.1)
TYPE:
|
min_keep_frac
|
Minimum fraction to keep (default 0.1)
TYPE:
|
max_keep_frac
|
Maximum fraction to keep (default 0.9)
TYPE:
|
next_gamma
The gamma this strategy would advance to, as a pure function of its inputs.
Kept separate from :meth:resample because gamma is state carried BETWEEN resamples: doing
the update inside a traced call would either capture a tracer in self or be dropped.
Returns a JAX scalar; the caller decides when to make it concrete.
resample
resample(points: ndarray, residuals: ndarray, domain: Any, tag: str, epoch: int, rng_key: ndarray, candidates: Optional[ndarray] = None) -> jnp.ndarray
Resample using causal time gating.
| PARAMETER | DESCRIPTION |
|---|---|
points
|
Current points (N, D)
TYPE:
|
residuals
|
Residual magnitudes (N,) or (B, N)
TYPE:
|
domain
|
domain object
TYPE:
|
tag
|
domain tag
TYPE:
|
epoch
|
Current epoch
TYPE:
|
rng_key
|
JAX random key
TYPE:
|
candidates
|
Pre-drawn candidate points from draw_candidates()
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
New points (N, D) |
R3
R3(resample_every: int = 100, resample_fraction: float = 0.7, start_epoch: int = 1000, threshold_mode: str = 'mean', min_keep_frac: float = 0.3, max_keep_frac: float = 0.9)
Residual-based Refinement and Resampling (R3).
Keeps points with high residuals and replaces low-residual points with new samples from the candidate pool. This focuses collocation on regions where the PDE is poorly satisfied.
Reference
"Residual-based adaptive sampling for physics-informed neural networks"
Initialize R3 resampling.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Not used for R3 (uses threshold instead)
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
threshold_mode
|
Threshold for keeping points ("mean", "median", or float)
TYPE:
|
min_keep_frac
|
Minimum fraction of points to keep (default 0.3)
TYPE:
|
max_keep_frac
|
Maximum fraction of points to keep (default 0.9)
TYPE:
|
resample
resample(points: ndarray, residuals: ndarray, domain: Any, tag: str, epoch: int, rng_key: ndarray, candidates: Optional[ndarray] = None) -> jnp.ndarray
Keep high-residual points, resample low-residual points.
| PARAMETER | DESCRIPTION |
|---|---|
points
|
Current points (N, D)
TYPE:
|
residuals
|
Residual magnitudes (N,) or (B, N)
TYPE:
|
domain
|
domain object
TYPE:
|
tag
|
domain tag
TYPE:
|
epoch
|
Current epoch
TYPE:
|
rng_key
|
JAX random key
TYPE:
|
candidates
|
Pre-drawn candidate points from draw_candidates()
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
New points (N, D) |
PINNFluence
PINNFluence(resample_every: int = 500, resample_fraction: float = 0.2, start_epoch: int = 2000, alpha: float = 1.0, c: float = 1.0, candidate_factor: float = 3.0)
PINNFluence: Influence function-based adaptive sampling (simplified).
Uses influence functions to score candidate points based on their potential impact on reducing the loss. Points with high influence scores are more likely to be sampled.
Note: This is a simplified version that uses gradient-based scoring rather than full influence functions for computational efficiency.
Reference
"Training Physics-Informed Neural Networks with Optimal Test Points"
Initialize PINNFluence resampling.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs (use larger values, ~500+)
TYPE:
|
resample_fraction
|
Fraction of points to replace
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
alpha
|
Score exponent for sampling (default 1.0)
TYPE:
|
c
|
Additive smoothing constant (default 1.0)
TYPE:
|
candidate_factor
|
Pool size multiplier (default 3.0)
TYPE:
|
resample
resample(points: ndarray, residuals: ndarray, domain: Any, tag: str, epoch: int, rng_key: ndarray, candidates: Optional[ndarray] = None) -> jnp.ndarray
Resample using simplified influence-based scoring.
| PARAMETER | DESCRIPTION |
|---|---|
points
|
Current points (N, D)
TYPE:
|
residuals
|
Residual magnitudes (N,) or (B, N)
TYPE:
|
domain
|
domain object
TYPE:
|
tag
|
domain tag
TYPE:
|
epoch
|
Current epoch
TYPE:
|
rng_key
|
JAX random key
TYPE:
|
candidates
|
Pre-drawn candidate points from draw_candidates()
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
New points (N, D) |
HA
HA(resample_every: int = 100, resample_fraction: float = 0.5, start_epoch: int = 1000, alternate: bool = True, random_first: bool = True)
Hybrid Adaptive (HA) resampling strategy.
Alternates between two phases: - Random: Fully random refresh of interior points - Adaptive: Retain high-residual points, fill remainder randomly
This provides regularization through random phases while still focusing on high-error regions during adaptive phases.
Reference
"Hybrid adaptive sampling for physics-informed neural networks" https://link.springer.com/article/10.1007/s10489-024-06195-2
Initialize HA resampling.
| PARAMETER | DESCRIPTION |
|---|---|
resample_every
|
Resample every N epochs
TYPE:
|
resample_fraction
|
Fraction of points to retain in adaptive phase (beta)
TYPE:
|
start_epoch
|
Start resampling after this many epochs
TYPE:
|
alternate
|
Whether to alternate between random and adaptive phases
TYPE:
|
random_first
|
Start with random phase if True
TYPE:
|
resample
resample(points: ndarray, residuals: ndarray, domain: Any, tag: str, epoch: int, rng_key: ndarray, candidates: Optional[ndarray] = None) -> jnp.ndarray
Hybrid adaptive resampling.
| PARAMETER | DESCRIPTION |
|---|---|
points
|
Current points (N, D)
TYPE:
|
residuals
|
Residual magnitudes (N,) or (B, N)
TYPE:
|
domain
|
domain object
TYPE:
|
tag
|
domain tag
TYPE:
|
epoch
|
Current epoch
TYPE:
|
rng_key
|
JAX random key
TYPE:
|
candidates
|
Pre-drawn candidate points from draw_candidates()
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
New points (N, D) |
RandomResampling
Randomly resample a fraction of collocation points.
Useful as a baseline to prevent overfitting to specific points.
resample
resample(points: ndarray, residuals: ndarray, domain: Any, tag: str, epoch: int, rng_key: ndarray, candidates: Optional[ndarray] = None) -> jnp.ndarray
Randomly replace a fraction of points with new samples from domain.
| PARAMETER | DESCRIPTION |
|---|---|
points
|
Current points (N, D)
TYPE:
|
residuals
|
Unused for random resampling
TYPE:
|
domain
|
domain object for sampling
TYPE:
|
tag
|
domain tag
TYPE:
|
epoch
|
Current epoch
TYPE:
|
rng_key
|
JAX random key
TYPE:
|
candidates
|
Pre-drawn candidate points from draw_candidates()
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ndarray
|
New points with some randomly replaced |
Parameter-efficient fine-tuning (jno.lora)
Attached with .lora(...) on a wrapped model — see LoRA.
jno.lora
Public namespace for LoRA adapter classes.
Import via::
import jno
net.lora(rank=4, wrapper=jno.lora.rsLoRALinear)
or::
from jno.lora import rsLoRALinear, DoRALinear, LoRAConv
ConvLike
Structural type for any conv module LoRA can wrap (excludes ConvTranspose).
DoRAConv
Weight-Decomposed Low-Rank Adaptation for conv layers.
Operates on the flattened weight (out_channels, flat_in) exactly
as DoRALinear does on a 2-D linear weight.
Reference: https://arxiv.org/abs/2402.09353
DoRALinear
Weight-Decomposed Low-Rank Adaptation (DoRA).
Decomposes the base weight into a magnitude vector and a unit-norm direction matrix. Low-rank matrices update the direction; the per-output magnitude is a small additional trainable vector.
W' = m * (W_0 + B @ A) / ||W_0 + B @ A||_row
Assumes the base layer computes x @ weight.T (+ bias).
Reference: https://arxiv.org/abs/2402.09353
IA3Conv
IA³ output-scaling adapter for conv layers.
Applies a learned per-output-channel scale to the conv weight before
the convolution. No low-rank matrices — trainable params = out_channels.
Reference: https://arxiv.org/abs/2205.05638
IA3Linear
Infused Adapter by Inhibiting and Amplifying Inner Activations (IA³).
Applies a single per-output learned scale vector to the linear output
(before adding bias). No low-rank matrices — trainable params per layer
is exactly out_features.
At initialisation scale = ones(out_features) so the output is
unchanged. rank is stored for API compatibility but unused.
y = (x @ weight.T) * scale + bias
Reference: https://arxiv.org/abs/2205.05638
LinearLike
Structural type for any linear module LoRA can wrap.
LoKrConv
Kronecker-product adapter for conv layers.
The weight delta is a Kronecker product of two small matrices, cropped to
(out_channels, flat_in) and reshaped to the original weight shape.
Reference: https://arxiv.org/abs/2212.10650
LoKrLinear
Kronecker-product Adapter (LoKr).
Parameterises the weight delta as a Kronecker product of two small matrices::
delta_W = kron(kron_A, kron_B)[:out_features, :in_features] * (alpha / rank)
kron_A has shape (rank, rank); kron_B has shape
(⌈out/rank⌉, ⌈in/rank⌉). The Kronecker product always covers the
target weight shape and is cropped. kron_B is initialised to zero so
the output is unchanged at init.
Reference: https://arxiv.org/abs/2212.10650
LoRAConv
Standard LoRA for convolutional layers.
Flattens the weight tensor to (out_channels, flat_in) and applies
a low-rank update identical to LoRALinear:
W' = W + (alpha/rank) * lora_B @ lora_A (in flattened space)
The updated weight is reshaped back and passed to the base conv via
eqx.tree_at, which XLA fuses into a single conv kernel call.
flat_in = in_channels // groups * prod(kernel_size)
LoRAFAConv
Frozen-A LoRA for conv layers. Only lora_B is trained.
Reference: https://arxiv.org/abs/2308.03303
LoRAFALinear
Frozen-A LoRA (LoRAFA).
Identical to LoRALinear but only lora_B is trained; lora_A
is frozen after random initialisation. Halves the number of trainable
adapter parameters at a small accuracy cost.
Reference: https://arxiv.org/abs/2308.03303
LoRALinear
Linear layer with frozen base weights and trainable LoRA adapters.
| ATTRIBUTE | DESCRIPTION |
|---|---|
base |
Original linear module (frozen during training). Accepts any
linear-like
TYPE:
|
lora_A |
Down-projection, shape
TYPE:
|
lora_B |
Up-projection, shape
TYPE:
|
rank |
LoRA rank (static).
TYPE:
|
alpha |
Scaling factor (static).
TYPE:
|
LoRAWrapper
Base class for LoRA-style parameter-efficient adapters.
Subclass this to support LoRA on any layer type. Three things required:
- Set
adapter_fieldsto the names of the trainable adapter attributes. - Implement
applies_to(leaf)— returnTruefor leaves this wrapper can handle. - Implement
__init__(base, rank, alpha, *, key)and__call__. - Implement
merge()— collapse adapters back into the base layer.
LoRAXSConv
LoRA-XS extra-small adapter for conv layers.
lora_A and lora_B are frozen (from SVD of flattened weight).
Only the small R (rank × rank) matrix is trained.
Reference: https://arxiv.org/abs/2405.17604
LoRAXSLinear
LoRA-XS: extra-small LoRA with a trainable r×r core matrix.
lora_A (r × in) and lora_B (out × r) are initialised from the
top singular vectors of the pretrained weight and then frozen. Only
the small R matrix (r × r, initialised to zero) is trained:
delta_W = B @ R @ A * (alpha / rank)
This can give better parameter efficiency than vanilla LoRA for the same
rank because R has r² parameters instead of r*(in+out).
Reference: https://arxiv.org/abs/2405.17604
MiLoRAConv
Minor Singular Values and Singular Vectors Adaptation for conv layers.
Like PiSSAConv but adapts the minor (least-significant) singular
components; base retains the principal directions.
Reference: https://arxiv.org/abs/2405.09913
MiLoRALinear
Minor Singular Values and Singular Vectors Adaptation (MiLoRA).
Like PiSSALinear but initialises the adapters from the least significant
singular components of the pretrained weight (the "noise" subspace). The frozen
base retains the principal (high-energy) components, so fine-tuning only updates
directions the base model treats as noise — preserving pretrained knowledge.
Reference: https://arxiv.org/abs/2405.09913
OFTConv
Orthogonal Fine-Tuning for conv layers.
The block-diagonal Cayley map is applied to the output-channel axis of
the flattened weight (out_channels, flat_in), identical to
OFTLinear on a 2-D weight.
Reference: https://arxiv.org/abs/2306.07280
OFTLinear
Orthogonal Fine-Tuning (OFT).
Fine-tunes the weight by multiplying with a block-diagonal orthogonal matrix
R. Each (rank × rank) block is parameterised via the Cayley map of a
learnable skew-symmetric matrix Q_i::
R_i = (I - Q_i)(I + Q_i)^{-1} (orthogonal for any skew-symmetric Q_i)
The full R is block-diagonal, assembled from n_blocks = ⌈out/rank⌉
blocks and cropped to (out_features, out_features). At init Q = 0
so R = I and the model output is unchanged.
alpha is stored for API compatibility; orthogonality is its own
regularisation and no explicit scaling is applied.
Reference: https://arxiv.org/abs/2306.07280
PiSSAConv
Principal Singular Values and Singular Vectors Adaptation for conv layers.
SVD is computed on the flattened weight (out_channels, flat_in).
Adapters start from the top-r singular components; base holds the residual.
Reference: https://arxiv.org/abs/2404.02948
PiSSALinear
Principal Singular Values and Singular Vectors Adaptation (PiSSA).
Initialises the adapter from the top-r singular components of the pretrained weight. The frozen base stores only the residual (the remaining singular components), so the model output is identical to the original at initialisation. Adapters therefore start from the most important weight directions rather than random noise, leading to faster convergence when fine-tuning pretrained models.
Reference: https://arxiv.org/abs/2404.02948
VeRAConv
Vector-based Random Matrix Adaptation for conv layers.
A and B are generated on-the-fly from a stored seed (XLA constants).
Only per-channel b (out_channels,) and d (rank,) are trained.
Reference: https://arxiv.org/abs/2310.11454
VeRALinear
Vector-based Random Matrix Adaptation (VeRA).
Replaces trainable A and B matrices with frozen random matrices that are generated on-the-fly from a stored integer seed and never materialised as Python/JAX arrays — JAX/XLA constant-folds their generation into the compiled kernel. Only two small per-layer scaling vectors are trained:
b(out_features,) — scales each output row of Bd(rank,) — scales each row of A
Trainable params per layer: out + rank (vs r·(in + out) for LoRALinear).
Checkpoints store only the seed and scaling vectors, not A or B.
Layers with the same seed and the same shape share A, B at the XLA level
via constant-expression deduplication (CSE), achieving true device-level sharing.
Reference: https://arxiv.org/abs/2310.11454
rsLoRAConv
Rank-Stabilized LoRA for conv layers. Scales by alpha/sqrt(rank).
Reference: https://arxiv.org/abs/2312.03732
rsLoRALinear
Rank-Stabilized LoRA (rsLoRA).
Replaces the alpha/rank scalar with alpha/sqrt(rank). This keeps
the effective gradient scale constant as rank grows, so larger ranks are
actually useful rather than numerically unstable.
Drop-in replacement for LoRALinear — identical API and parameter count.
Reference: https://arxiv.org/abs/2312.03732
apply_lora
apply_lora(model: Module, rank: int = 0, alpha: float = 1.0, *, key: Array, target: Optional[str] = None, specs: Sequence[LoRASpec] | None = None, wrappers: type[LoRAWrapper] | Sequence[type[LoRAWrapper]] | None = None, param_mask: Any = None) -> eqx.Module
Apply LoRA adapters to matching layers in model.
Two calling conventions:
-
Uniform:
apply_lora(model, rank, alpha, key=key)Wraps all matching layers with the same rank/alpha. Passtargetto filter by pytree-path regex, andwrappersto override the adapter class (default:LoRALinear). -
Per-target:
apply_lora(model, key=key, specs=[...])Each spec{"target": regex, "rank": int, "alpha": float}may also carry an optional"wrapper"key to use a different adapter class for that group. The first matching spec wins.
| PARAMETER | DESCRIPTION |
|---|---|
param_mask
|
Optional boolean pytree mirroring the model structure.
When provided, only modules whose corresponding sub-tree contains
at least one
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Module
|
A new model with |
lora_trainable_filter
Return a filter-spec pytree that marks only LoRA adapter arrays as trainable.
Works with any LoRAWrapper subclass via its adapter_fields declaration.
merge_lora
Collapse all LoRA adapters back into their base layers.
partial_lora_trainable_filter
Return a filter-spec for partial LoRA (mask(M).lora() without freeze()).
Three-bucket logic:
- LoRA adapter arrays (b, d, lora_A, lora_B, …) → True (trained)
- Base arrays inside LoRA wrappers → False (frozen)
- Arrays outside any LoRA wrapper → True (trained)
Use this when the user says mask(M).lora() without a prior freeze():
selected layers get LoRA with frozen bases; everything else stays trainable.
Training trackers (jno.trackers)
Diagnostics attachable as callbacks — see Explainability.
jno.trackers
Live training-loop trackers for adaptive components and diagnostics.
The objects produced by this module are the same callback classes exposed
through :class:jno.callbacks for explainability, but surfaced under a
namespace that matches their conceptual role: trackers whose latest value
is read by other in-loop components (e.g. NTK-balanced loss weighting,
gradient-norm-balanced loss weighting).
Each tracker is registered via the solver's callbacks=[...] list and
exposes two attributes on top of the existing cb.result history:
- :attr:
tracker.value <jno.utils.adaptive.callbacks._LiveValue.value>— the most recent computed value as a numpydict(Noneuntil the first interval fires). - :attr:
tracker.latest_epoch <jno.utils.adaptive.callbacks._LiveValue.latest_epoch>— epoch index when :attr:valuewas last updated.
Note
:func:jno.tracker (singular) is unrelated: it wraps a single
:class:~jno.trace.Placeholder expression and stores it in
logs["track_stats"] (scalar → 2-D float array; non-scalar →
list[list[np.ndarray]]). The explainability trackers here cannot use
that path because they depend on the compiled-constraints function,
partition pytrees, masks, and host-side reductions (Lanczos, eigh).
See :mod:jno.utils.adaptive.weights for adaptive weight schemes
(NTKBalanced, GradientNormBalanced) that consume these trackers.
CosSimilarityCallback
Track pairwise gradient cosine similarity between loss terms.
Computes the full (N × N) cosine similarity matrix (upper triangle
carries the pairwise values; diagonal is always 1).
| PARAMETER | DESCRIPTION |
|---|---|
interval
|
Compute every n outer training steps. Default
TYPE:
|
mask
|
Optional pytree of booleans — see :class:
TYPE:
|
Example::
cb = jno.callbacks.cos_similarity(interval=100)
crux.solve(5000, callbacks=[cb])
print(cb.result["cos_sim_matrix"]) # (n_samples, N, N)
GradientAlignmentCallback
Track the total gradient alignment scalar during training.
Computes 2‖(1/N) Σ ĝᵢ‖² − 1 with ĝᵢ = gᵢ / ‖gᵢ‖ (Eq. 3.1,
[2502.00604]), a value in [-1, 1] that measures how well all loss
gradients point in the same direction. Near 1 means perfect
alignment; 0 means orthogonal; near -1 means anti-aligned
(gradients actively cancel). For N = 2 this reduces to the
ordinary cosine similarity (Proposition 1 of the paper).
| PARAMETER | DESCRIPTION |
|---|---|
interval
|
Compute every n outer training steps. Default
TYPE:
|
mask
|
Optional pytree of booleans — see :class:
TYPE:
|
Example::
cb = jno.callbacks.gradient_alignment(interval=100)
crux.solve(5000, callbacks=[cb])
print(cb.result["alignment"]) # (n_samples,)
GradientNormsCallback
Track the gradient norm of each loss term during training.
At every interval outer steps, computes ‖∇L_i‖₂ for each
constraint i by differentiating through the constraint function.
| PARAMETER | DESCRIPTION |
|---|---|
interval
|
Compute every n outer training steps. Default
TYPE:
|
mask
|
Optional pytree of booleans matching
TYPE:
|
Example::
cb = jno.callbacks.gradient_norms(interval=100)
crux.solve(5000, callbacks=[cb])
print(cb.result["norms"]) # (n_samples, n_constraints)
HessianSpectrumCallback
HessianSpectrumCallback(k: int = 10, n_iter: int = 30, interval: int = 500, mask: Any = None, constraints: Any = None)
Track the top-k Hessian eigenvalues of the total training loss.
Constructs :math:\nabla^2_\theta L implicitly via Hessian-vector
products (jvp(grad(L), …, …)) and runs n_iter Lanczos
iterations with full reorthogonalisation. The resulting tridiagonal
is eigendecomposed host-side via scipy.linalg.eigh_tridiagonal
(Sec. 3.1-3.2, [1912.07145]) to obtain the top-k eigenvalues.
The largest eigenvalue is the sharpness of the loss surface at the current iterate (Sec. 2.2, [1609.04836]); high values predict a sharp minimum, typically associated with worse generalisation.
| PARAMETER | DESCRIPTION |
|---|---|
k
|
Number of largest eigenvalues to report. Default
TYPE:
|
n_iter
|
Number of Lanczos iterations. Default
TYPE:
|
interval
|
Compute every n outer training steps. Default
TYPE:
|
mask
|
Optional pytree of booleans matching the trainable structure. Essential for large models.
TYPE:
|
constraints
|
Optional list of constraint placeholders to scope the
Hessian to the mean of those constraints' losses (instead of
the total training loss). Pass the same Python objects given
to
TYPE:
|
InputSensitivityCallback
Evaluate any placeholder expression at training collocation points.
Most useful for input-saliency: pass u.d(x) (single coordinate) or
jno.Jacobian(u, [x, y]) (multi-coordinate) and the callback will
record the per-point sensitivity of the network output to its inputs
every interval outer steps. Conceptually equivalent to the
class-saliency map of Sec. 3 of [1312.6034] — high magnitude points
correspond to spatial regions where the network response is most
responsive to a perturbation of the input coordinate.
The expression is compiled via
:meth:~jno.trace_compiler.TraceCompiler.compile_multi_expression
(the same machinery used by the solver's constraints and trackers),
so any composite placeholder expression works — not just first
derivatives.
| PARAMETER | DESCRIPTION |
|---|---|
expr
|
A :class:
TYPE:
|
interval
|
Compute every n outer training steps. Default
TYPE:
|
Example::
cb = jno.callbacks.input_sensitivity(u.d(x), interval=100)
crux.solve(5000, callbacks=[cb])
print(cb.result["values"].shape) # (n_samples, *expr_shape)
LossLandscapeCallback
LossLandscapeCallback(interval: int = 500, mask: Any = None, n_grid: int = 15, alpha_range: float = 1.0)
Evaluate the 2-D loss landscape periodically during training.
At every interval outer steps, samples two random filter-normalized
directions in parameter space and evaluates the total loss on an
(n_grid × n_grid) perturbation grid centred on the current
parameters. The directions change each call (stochastic sampling).
This is expensive: each call requires n_grid² forward passes.
Choose a large interval (e.g. 500–1000) for typical runs.
| PARAMETER | DESCRIPTION |
|---|---|
interval
|
Compute every n outer training steps. Default
TYPE:
|
mask
|
Optional pytree of booleans matching
TYPE:
|
n_grid
|
Number of grid points per axis. Default
TYPE:
|
alpha_range
|
Perturbation scale in units of
TYPE:
|
Example::
cb = jno.callbacks.loss_landscape(interval=500, n_grid=15)
crux.solve(5000, callbacks=[cb])
landscapes = cb.result["landscapes"] # (n_samples, 15, 15)
NTKSpectrumCallback
Track the empirical Neural Tangent Kernel eigenvalue spectrum.
For a :class:~jno.trace.NetworkGradient placeholder u.grad(net)
(per-point parameter Jacobian :math:J \in \mathbb{R}^{N \times P})
constructs the empirical NTK :math:K = J J^\top at every
interval outer steps, subsampled to n_points collocation
points, and reports the top-k eigenvalues, condition number, and full
spectrum.
The eigenvalue spread is the canonical diagnostic for PINN spectral bias (Sec. 3-4, [2007.14527]): widely separated eigenvalues mean some directions in parameter space train orders of magnitude faster than others, producing the characteristic PINN failure mode where high-frequency features lag low-frequency ones.
| PARAMETER | DESCRIPTION |
|---|---|
grad_expr
|
A :class:
TYPE:
|
n_points
|
Number of collocation points to subsample for the
kernel. Cost is
TYPE:
|
top_k
|
Number of largest eigenvalues to report. Default
TYPE:
|
interval
|
Compute every n outer training steps. Default
TYPE:
|
Example::
cb = jno.callbacks.ntk_spectrum(u.grad(u_net), n_points=128, top_k=10)
crux.solve(10_000, callbacks=[cb])
print(cb.result["lambda_max"]) # (n_samples,)
print(cb.result["condition_number"]) # (n_samples,)
ResidualStatsCallback
Track per-constraint residual distributions during training.
At every interval outer steps, evaluates each constraint's
un-reduced residual array (shape (B, T, ...)) and records four
summary statistics per constraint — mean, std, max, and 99th percentile
— plus a histogram of the raw residual magnitudes (Sec. 3, [2207.10289]).
A constraint whose max or p99 stays orders of magnitude above
the others indicates a region of the domain where the PDE is poorly
satisfied; this complements :class:GradientNormsCallback which only
reflects a constraint's aggregated contribution to the gradient.
| PARAMETER | DESCRIPTION |
|---|---|
interval
|
Compute every n outer training steps. Default
TYPE:
|
constraints
|
Optional list of constraint placeholders to scope the
callback to a subset of the solver's constraints. Pass the
same Python objects that were given to
TYPE:
|
Example::
pde_loss = pde.mse # assign once
bc_loss = bc.mse
solver = jno.core([pde_loss, bc_loss])
cb = jno.callbacks.residual_stats(interval=100, constraints=[pde_loss])
crux.solve(5000, callbacks=[cb])
print(cb.result["maxes"]) # (n_samples, 1) — just pde_loss
result
property
Collected data as numpy arrays.
Keys
epochs — (S,) int array of sampled outer steps.
means — (S, K) float32 per-constraint residual mean,
where K = len(constraints) when a subset was selected,
else N_constraints.
stds — (S, K) float32 per-constraint residual std.
maxes — (S, K) float32 per-constraint residual max.
p99 — (S, K) float32 per-constraint 99th-percentile.
indices — (K,) int array of solver-side constraint indices
in the same order as the columns of means/stds/etc.
gradient_norms
Live tracker for per-loss gradient L2 norms.
See :class:~jno.utils.adaptive.callbacks.GradientNormsCallback for the
full reference. After the first interval, tracker.value["norms"]
is a (N,) numpy array of per-constraint gradient norms.
cos_similarity
Live tracker for pairwise gradient cosine similarity.
After the first interval, tracker.value["cos_sim_matrix"] is the
(N, N) matrix.
gradient_alignment
Live tracker for the total gradient alignment scalar.
After the first interval, tracker.value["alignment"] is a Python
float in [-1, 1].
residual_stats
Live tracker for per-constraint residual mean/std/max/p99.
After the first interval, tracker.value contains means, stds,
maxes, p99 (each (K,)) and indices ((K,) int).
ntk_spectrum
ntk_spectrum(grad_expr, n_points: int = 256, top_k: int = 10, interval: int = 500) -> NTKSpectrumCallback
Live tracker for the empirical NTK eigenvalue spectrum.
After the first interval, tracker.value contains eigvals_topk,
lambda_min, lambda_max, condition_number, all_eigvals,
and trace (sum of eigenvalues — convenient for NTK loss balancing).
hessian_spectrum
hessian_spectrum(k: int = 10, n_iter: int = 30, interval: int = 500, mask=None, constraints=None) -> HessianSpectrumCallback
Live tracker for the top-k Hessian eigenvalues (sharpness).
After the first interval, tracker.value contains eigvals and
sharpness.
input_sensitivity
Live tracker for an arbitrary placeholder expression.
After the first interval, tracker.value["values"] is the evaluated
array (e.g. u.d(x) per collocation point).
Deployment (jno.iree)
jno.utils.iree.IREEModel
IREEModel(vmfb_bytes: bytes, module_name: str, device: str = 'local-sync', input_dtypes: tuple[str, ...] | None = None)
Compiled IREE model — serialisable via jno.save / jno.load.
Typical usage::
# Compile once
model = IREEModel.compile(my_jax_fn, sample_inputs)
# Save / load using the standard jno API
jno.save(model, "model.pkl")
model = jno.load("model.pkl")
output = model(x, y)
The IREE runtime objects (context, module handle) are not picklable
and are reconstructed transparently on unpickle via __setstate__.
compile
classmethod
compile(func: Callable, sample_inputs: tuple, target_backend: str = 'llvm-cpu', optimization_level: int = 3) -> IREEModel
Compile a JAX function to an :class:IREEModel.
The returned object can be persisted with jno.save / jno.load::
model = IREEModel.compile(infer, sample_inputs)
jno.save(model, "model.pkl")
| PARAMETER | DESCRIPTION |
|---|---|
func
|
JAX-compatible function to compile.
TYPE:
|
sample_inputs
|
Example inputs (used for shape inference).
TYPE:
|
target_backend
|
IREE target backend (default
TYPE:
|
optimization_level
|
IREE optimisation level 0–3 (default
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
IREEModel
|
A ready-to-use :class: |