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 (
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): No parallelism, single device (default) - (2, 1): Pure data parallelism on 2 GPUs - 2x throughput - (1, 2): Pure model parallelism on 2 GPUs - fit 2x larger models - (4, 1): Data parallelism on 4 GPUs - 4x throughput - (2, 2): Hybrid parallelism on 4 GPUs - 2x data, 2x model - (4, 2): Hybrid parallelism on 8 GPUs - 4x data, 2x model Note: batch * model must equal the total number of available devices. 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), automatically expanded to (n_devices, 1) for pure data parallelism when multiple devices are available.
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.
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) |
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.
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) -> 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) -> 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).
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; lstsq (LSQR, least-squares); chebyshev (polynomial) |
| Linear — multigrid | amg (GPU AMG / NVIDIA AmgX via jaxamg) |
| Nonlinear | newton, picard |
| 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), exponential (exponential integrator), adaptive (step-doubling adaptive step size) |
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 block
subspace iteration (Ericsson & Ruhe, Math. Comp. 35, 1980; Bathe & Wilson, IJNME 6, 1973):
θ = 1/(λ−σ) makes the near-σ modes dominant with enormous transformed gaps, so the transformation is
its own preconditioner and precond= is rejected there. 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 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.
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.
Preconditioners (jno.precond, for the iterative solvers): jacobi, chebyshev,
nystrom (randomized low-rank — the rung between jacobi and multigrid),
amg (algebraic multigrid), gmg (geometric multigrid — a structured-grid V-cycle),
ams (H(curl) auxiliary-space Maxwell), form (weak-form auxiliary operator),
inner (any solver as M⁻¹), block_diag / triangular (block / Schur), and cached.
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-Krylov
Defaults when a slot is None are unchanged from the historic behaviour: Jacobi-preconditioned
BiCGStab (steady linear) and Jacobian-free Newton-Krylov (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)
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
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; no vmap batching rule upstream (trait
vmap="no") -- use a Krylov solver inside vmapped/batched solves.
| 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:
|
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, and the only shipped direct solver with a native vmap batching rule. 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.
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 = False) -> NonlinearSolver
Newton root-find, as a configurable slot. Two inner-solve modes:
- default (matrix-free) --
J @ vfrom a JVP, inner matrix-free solve (default BiCGStab, or thelinear=slot), implicit differentiation vialax.custom_root. The historic behaviour. direct=True(sparse-direct) -- factorize the ASSEMBLED tangent each step with a sparse LU instead of an iterative inner solve. Robust on indefinite / ill-conditioned systems -- a Taylor-Hood velocity/pressure saddle, a stiff Carman-Kozeny phase-change drag -- where the matrix-free BiCGStab has no saddle-point preconditioner and stalls. Still differentiable (implicit diff with a direct, transposable tangent solve at the root). Composes only where the assembler provides the tangent:fem.solve(nonlinear=jno.solve.newton(direct=True))on a native nonlinear problem (steady or the transient stepper); thelinear=/precond=slots are then unused.
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.
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) -> 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.
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) -> 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. Each field is solved on
its own — solving a GROUP of fields together (a Stokes velocity/pressure pair inside one sweep) is
not wired.
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 — or a shift
lands on an eigenvalue and the inner factorization degenerates — the result is NaN-poisoned
rather than silently under-converged — 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). 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.
remesh
remesh(*, anisotropic: bool = False, max_dofs: int | None = None, every: int = 5, 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.
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 at a constant
budget and carries the state across (basis-aware transfer), so the mesh tracks a moving feature and
coarsens its wake instead of ratcheting up::
fem.solve(adapt=jno.solve.remesh(anisotropic=True, max_dofs=6000, every=4))
anisotropic=True refines on a Hessian metric (stretched elements aligned to the solution's
curvature) 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. Metric-based DOF control is approximate, so max_dofs is honoured only
loosely in that mode.
Steady-only: max_iters, tol, eps (a relative-change plateau detector, not a certified
bound). Transient-only: every, metric_field.
| PARAMETER | DESCRIPTION |
|---|---|
anisotropic
|
Hessian-metric refinement instead of isotropic ZZ + Dörfler marking.
TYPE:
|
max_dofs
|
Vertex budget. Steady: stop once reached. Transient: the constant target.
TYPE:
|
every
|
Transient only — remesh every
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.
relocate
relocate(*, method: str = 'descent', max_iters: int = 8, lr: float = 0.003, quality_floor: float = 0.1, relax: int = 60, relax_step: float = 0.1) -> 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 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-03where the full map stayed positive throughout. Freeing tangential axes recovers nearly all of it. Either method checksdet Jeach 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 fromrelax=40to300atrelax_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:
remeshwhen 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.
| PARAMETER | DESCRIPTION |
|---|---|
method
|
TYPE:
|
max_iters
|
Outer relocation rounds.
TYPE:
|
lr
|
TYPE:
|
quality_floor
|
TYPE:
|
relax
|
TYPE:
|
relax_step
|
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
AdaptSpec
|
The adaptation spec to pass as
TYPE:
|
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).
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.
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.
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); :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 (from fem.offsets), 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).
grid
property
Structured-grid descriptor {shape, spacing, origin} when the operator lives on a regular
grid (jno.domain(..., structured=True)), 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"]).
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.
gmg
Geometric-multigrid V-cycle preconditioner for a structured grid (jno.domain(...,
structured=True)).
Builds a coarsen-by-2 grid hierarchy and applies one V-cycle as M⁻¹: damped-Jacobi smoothing
(n_pre/n_post sweeps, omega damping — default the model-problem optimum 2d/(2d+1)),
full-weighting restriction, multilinear prolongation, rediscretised coarse Laplacians, and a
dense solve at the coarsest level (stops coarsening below min_size nodes/axis or at an odd cell
count). Convergence is grid-independent — ~0.1 residual reduction per V-cycle, O(N) work — on
Poisson / Helmholtz-type operators. Matrix-free and differentiable; the V-cycle is a fixed linear
operator, so standard GMRES (not FGMRES) suffices.
Use it as fem.solve(linear=jno.solve.gmres(), precond=jno.precond.gmg()) on a structured domain;
a structured jno.fdm solve already uses it automatically. Raises if the operator has no
structured grid, or the grid is too small to coarsen. v1 is constant-coefficient (the rediscretised
coarse operator); a Galerkin RAP coarse operator for variable coefficients is future work.
Reference: A. Brandt, Multi-Level Adaptive Solutions to Boundary-Value Problems, Mathematics of Computation 31(138), 1977.
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:
- 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.
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. With an iterative inner, drive the outer solve with
jno.solve.fgmres() (flexible preconditioning).
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().
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.
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().
amg
amg(*, cycles: int = 1, max_levels: int = 10, coarse_size: int = 100, smoother_degree: int = 3) -> _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). 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.
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.
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).
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.
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
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.
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) -> _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).
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.