FD + FE backendsGPU / multi-GPUdifferentiableopen source

jaxmag — micromagnetics, differentiable.

A GPU micromagnetic simulation package built on JAX: finite-difference and finite-element solvers, FFT and fast-multipole (jaxFMM) demagnetisation, adaptive Dopri5/BDF time integration with demag extrapolation — and gradients through the whole LLG for inverse design.

▶ Run a simulation in your browser Source on GitLab

Why jaxmag

⚡ Fast by construction

JIT-compiled XLA kernels, cheap/expensive field splitting with demag extrapolation: one demag evaluation per accepted step instead of six.

📐 FD and FE

Regular-grid FFT demag for films and boxes; tetrahedral P1 finite elements with jaxFMM fast-multipole demag for curved geometries — including surface-following anisotropy and interfacial DMI on measured topographies.

∇ Differentiable

Autodiff through energies, fields and full trajectories. Inverse problems — reconstructing magnetisation, fitting D and A, shape optimisation — become gradient descent.

🖧 Multi-GPU

Sharded FFT demag scales to 8× H100; a 216-million-cell permalloy cube (N=600) runs as a single simulation.

🌀 Topology-aware analysis

Built-in observables: skyrmion radius, Berg–Lüscher topological charge, trajectories, energy terms.

✅ Validated

muMAG SP4 and the Cortés-Ortuño DMI standard problem reproduced against mumax3, OOMMF and magnum.pi; FD DMI fields bit-exact vs magnum.np.

Benchmarks

ProblemjaxmagReference
muMAG SP4, 1 nm grid (187k cells), 1 ns switch 22.3 s (H100, Dopri5+extrap)mumax3: 59 s (A100)
Local FE field (exch+aniso+DMI), 109k-node mesh 1.05 ms / callmagnum.pi: 48 ms (46× slower)
Full 200 ps FE relaxation on measured topography 116 smagnum.pi: 430 s
Cortés-Ortuño DMI standard problem (R → 22 nm) 4.4 s (FD)OOMMF: 21.9 nm ✓
N=600 permalloy cube, 5M tets, 1 ns (FE+jaxFMM) 5.1 h (H100)FEMME BEM: 7.3 h (32 cores)

Quickstart

# install (pulls jax-ode automatically)
pip install git+https://gitlab.com/dietersuess1/jaxmag.git

# relax a Néel skyrmion (Cortés-Ortuño standard problem)
from jaxmag.fd import make_mesh, make_demag, make_exchange, make_dmi
from jaxmag.core import make_llg_equation
from jax_ode import make_dopri5_split

mesh  = make_mesh(lx=100e-9, ly=100e-9, lz=2e-9, dx=2e-9, dy=2e-9, dz=2e-9)
demag = make_demag(mesh, Ms=860e3)
cheap = lambda m: make_exchange(mesh, A=13e-12, Ms=860e3)(m) \
              + make_dmi(mesh, D=3e-3, Ms=860e3)(m)
step  = make_dopri5_split(cheap, demag, demag.kernels_tuple,
                          equation_fn=make_llg_equation(alpha=1.0),
                          shape=mesh.shape)

Download

Get the complete jaxmag source as a self-contained archive — including the vendored jax-ode and jaxfmm dependencies, all example scripts, and full Markdown documentation (README — API walkthrough, INSTALL — setup, EXAMPLES — guide to every example, TESTING — validation against muMAG standard problems).

⬇️ Download jaxmag-bundle.zip (0.3 MB)   or clone on GitLab →

# after unzipping:
pip install "jax[cuda12]"                       # or plain "jax" for CPU
pip install ./vendor/jax-ode ./vendor/jaxfmm .
python examples/sp4_fd.py                        # muMAG standard problem 4

Run online — describe it, we compute it

Type what you want to simulate (English or German). An AI assistant maps it onto a jaxmag template, runs it on a GPU, and returns a PDF report (parameters, final state, trajectory, topological charge) right here.

🧪 Standard problem 4

Permalloy switching dynamics, 500×125×3 nm — the classic code benchmark.

🌀 DMI skyrmion

Cortés-Ortuño standard problem: skyrmion relaxes to R ≈ 22 nm, Q = −1.

📦 Custom relaxation

Your own box, material and field — see where the magnetisation settles.

Jobs & reports