Installation

Requires Python 3.11–3.13.

PyPI

pip install jax-numerical-operators

That is the whole library: FEM, FDM, the solver stack, PINNs and SciML all work from the core install. The finite-element assembler's fenics-basix backend (imports as basix) is a core dependency, and jno.fdm has no optional dependency at all — neither needs an extra.

The core install is CPU-only on purpose — everything runs on a plain jax wheel, and a laptop, CI runner, or macOS machine should not pay for multi-GB CUDA wheels it cannot use. An NVIDIA GPU is one extra away:

pip install "jax-numerical-operators[cuda]"

Optional extras

Everything below is imported lazily: the core install works without it, and each missing backend raises a clear ImportError naming the extra that provides it. Install any of them as pip install "jax-numerical-operators[<extra>]"; extras combine as [fem,rcwa].

Extra	Enables	Pulls in
`[cuda]`	NVIDIA GPU support — swaps in the CUDA-capable JAX build (same version pin as the core `jax`)	`jax[cuda]`
`[fem]`	Everything `jno.fem` reaches for beyond the core — the one-liner. Meta-extra = `[mesh]` + `[pardiso]` + `[cudss]` + `[metis]`.	see the four below
`[fdm]`	What `jno.fdm` solves reach for beyond the core (finite differences themselves are core). Meta-extra = `[metis]`.	`pymetis`
`[mesh]`	Adaptive/anisotropic remeshing behind `fem.solve(adapt=...)`	`mmgpy`
`[pardiso]`	`jno.solve.lu(backend="pardiso")` — Intel MKL PARDISO, multithreaded CPU sparse-direct. The fastest factorization measured here, and the answer when a factorization exceeds GPU memory. x86-64 only.	`pypardiso`
`[cudss]`	`jno.solve.lu(backend="cudss")` — NVIDIA cuDSS, the fastest repeated solve (shift-invert eigensolves, constant-operator transients). Linux x86-64.	`nvmath-python`, `nvidia-cudss-cu12`, `cupy`
`[metis]`	The graph partition behind `jno.precond.schwarz()` (METIS)	`pymetis`
`[rcwa]`	The RCWA solver (`jno.rcwa`, periodic-layered electromagnetics)	`fmmax`
`[amg]`	GPU algebraic multigrid — `jno.solve.amg` / `jno.precond.ams`	`jaxamg` (builds a CUDA extension against a prebuilt AmgX 2.5+; needs that toolchain)
`[iree]`	`model.to_iree(...)` export	`iree-base-compiler`, `iree-base-runtime`
`[dev]`	Test/lint tooling for working on jNO itself	pytest, ruff-era tooling, …

pip install "jax-numerical-operators[fem]"        # FEM with remeshing + both direct backends
pip install "jax-numerical-operators[rcwa]"       # + the Fourier-modal EM solver
pip install "jax-numerical-operators[fem,rcwa]"   # combine freely

Finite differences ship in the core install; [fdm] only adds the solver back-ends an FDM solve can use (today: METIS for jno.precond.schwarz()).

The platform markers are part of the design: on a machine where a backend cannot run (PARDISO on arm64, cuDSS on macOS) the extra installs cleanly and simply lands without it — no resolver failure. A marker cannot detect whether a CUDA device is present, though: the cuDSS wheels install on any Linux x86-64 host and raise at call time without a GPU, which is why jno.solve.lu() never defaults to them.

The JAX pin is tight on purpose: jNO tracks a single JAX minor version per release to avoid silently breaking on JAX API changes, and the [cuda] extra carries the identical pin so CPU and GPU installs can never drift apart.

If you need a specific CUDA toolkit version instead of the [cuda] extra, install JAX from its own package index before installing jNO:

# CUDA 12 example
pip install --upgrade "jax[cuda12]>=0.10.1,<0.11"
pip install jax-numerical-operators

To pin a different JAX version locally for an experiment, see the JAX install matrix.

Clone + Pixi

For development or to run examples from source. Requires pixi.

git clone https://github.com/FhG-IISB/jNO.git
cd jNO
pixi install

Common tasks:

pixi run fmt     # format with ruff
pixi run lint    # lint and auto-fix
pixi run test    # run the test suite

The pip extras above map to pixi environments: pixi run -e fem ..., -e rcwa, -e iree, and -e dev (adds matplotlib and the test tooling).

Check what you actually got

Eleven optional extras means a missing one usually surfaces much later, as an ImportError in the middle of a solve. This reports it up front — devices, precision, and which backends imported:

import importlib.util as iu
import jax, jno

print(f"jNO {jno.__version__} · JAX {jax.__version__} · x64={jax.config.jax_enable_x64}")
print("devices:", [f"{d.platform}:{d.id}" for d in jax.devices()])

OPTIONAL = [
    ("mmgpy",        "[fem]",  "adaptive remeshing — fem.solve(adapt=...)"),
    ("pypardiso",    "[fem]",  'sparse-direct on CPU — jno.solve.lu(backend="pardiso")'),
    ("nvmath",       "[fem]",  'sparse-direct on GPU — jno.solve.lu(backend="cudss")'),
    ("jaxamg",       "[amg]",  "GPU algebraic multigrid — jno.solve.amg()"),
    ("fmmax",        "[rcwa]", "anisotropic RCWA layers"),
    ("iree.runtime", "[iree]", "ahead-of-time export — model.to_iree(...)"),
    ("pymetis",      "[metis]", "graph partition — jno.precond.schwarz()  ([fem] and [fdm] pull it in)"),
    ("matfree",      "",       "matrix functions — jno.solve.logdet / applyfun"),
    ("pyamg",        "",       "CPU algebraic multigrid — jno.precond.amg()"),
]
for mod, extra, unlocks in OPTIONAL:
    try:
        have = iu.find_spec(mod) is not None
    except ModuleNotFoundError:          # a dotted name whose parent package is absent
        have = False
    print(f"  {'yes' if have else ' no'}  {mod:<13} {extra:<7} {unlocks}")

In the repository's default pixi environment (which adds pymetis), on a CPU, it prints:

jNO 0.4.0 · JAX 0.10.1 · x64=False
devices: ['cpu:0']
   no  mmgpy         [fem]   adaptive remeshing — fem.solve(adapt=...)
   no  pypardiso     [fem]   sparse-direct on CPU — jno.solve.lu(backend="pardiso")
   no  nvmath        [fem]   sparse-direct on GPU — jno.solve.lu(backend="cudss")
   no  jaxamg        [amg]   GPU algebraic multigrid — jno.solve.amg()
   no  fmmax         [rcwa]  anisotropic RCWA layers
   no  iree.runtime  [iree]  ahead-of-time export — model.to_iree(...)
  yes  pymetis       [metis] graph partition — jno.precond.schwarz()  ([fem] and [fdm] pull it in)
  yes  matfree               matrix functions — jno.solve.logdet / applyfun
  yes  pyamg                 CPU algebraic multigrid — jno.precond.amg()

x64=False is the JAX default, and FEM needs True

Data precision is JAX's flag, not a jNO setting — jNO only respects it (the per-model .dtype() control is a separate thing). A FEM assembly accumulates in float64 and an unstructured-mesh solve will not reach its tolerance without it:

import jax
jax.config.update("jax_enable_x64", True)   # before you build a domain

Set it before creating arrays, domains or models: the flag affects only what is created after it, so flipping it mid-script leaves you with a mix.

Docker

CPU:

docker run --rm ghcr.io/fhg-iisb/jno:latest

GPU (requires NVIDIA drivers and NVIDIA Container Toolkit):

docker run --rm --gpus all ghcr.io/fhg-iisb/jno:latest

Build locally:

docker build -t jno:latest .
docker run --rm --gpus all jno:latest

What the image contains

The core install plus a CUDA-capable JAX — the image's pixi environment carries jax-cuda12-plugin[with-cuda], so --gpus all needs only the host driver and the NVIDIA Container Toolkit, no CUDA base image.

It does not carry the optional extras: no mmgpy (adaptive remeshing), no pypardiso / nvmath (accelerated direct solvers), no fmmax (anisotropic RCWA), no iree. Run the verification snippet inside the container if you are unsure what you have. Tags are published per release ({{version}} and latest), so latest tracks the most recent release, not main.