magpylib_jax¶
magpylib_jax
Analytic magnetic fields you can differentiate, compile, and optimize through.
magpylib_jax is a clean-room, JAX-native reimplementation of
Magpylib. It keeps Magpylib’s ergonomic object and
functional APIs and its analytic field models, but swaps the numerical core for a
differentiable, JIT-compilable, vmap-friendly one — so the same field computation you use for
analysis drops straight into a jax.grad optimization loop, on CPU, GPU, or TPU.
import jax
jax.config.update("jax_enable_x64", True) # float64, for magpylib parity
import magpylib_jax as mpj
src = mpj.magnet.Cuboid(polarization=(0, 0, 1.0), dimension=(1, 1, 1))
B = src.getB((2.0, 0.0, 0.0)) # tesla
dBz_dh = jax.grad(lambda h:
mpj.magnet.Cuboid(polarization=(0, 0, 1.0), dimension=(1, 1, h)).getB((2, 0, 0))[2])(1.0)
What you get¶
Exact grad / jacfwd / jacrev through getB/getH/getJ/getM and getFT, w.r.t. geometry,
pose, and excitation — no finite differences.
The field core runs under jax.jit / vmap / XLA on CPU, GPU, and TPU.
getFT gives force and torque by autodiff of m·B — eps-free and exact, unlike Magpylib’s
finite-difference approach.
Same source classes, Collection, Sensor, motion, and getB/…/getFT/show/mu_0. Often just swap
the import.
Follows your JAX config: float32 by default, float64 for bit-level parity. The library never
mutates the global config.
All 12 source families, ≥95% coverage, a full Magpylib parity suite, and derivations with citations.
Where to next¶
Install, compute your first field and gradient, and pick a precision.
Tutorials: fields, motion & paths, collections, visualization, force/torque, and optimization.
The physical conventions, per-source closed-form models, numerics, differentiability, and the papers behind each kernel.
The API reference and a clickable map from the public API down to the analytic kernels.