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 GitLabJIT-compiled XLA kernels, cheap/expensive field splitting with demag extrapolation: one demag evaluation per accepted step instead of six.
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.
Autodiff through energies, fields and full trajectories. Inverse problems — reconstructing magnetisation, fitting D and A, shape optimisation — become gradient descent.
Sharded FFT demag scales to 8× H100; a 216-million-cell permalloy cube (N=600) runs as a single simulation.
Built-in observables: skyrmion radius, Berg–Lüscher topological charge, trajectories, energy terms.
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.
| Problem | jaxmag | Reference |
|---|---|---|
| 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 / call | magnum.pi: 48 ms (46× slower) |
| Full 200 ps FE relaxation on measured topography | 116 s | magnum.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) |
# 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)
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
Type what you want to simulate (English or German). An AI assistant maps it onto a jaxmag template, runs it on a cloud GPU, and returns a PDF report (parameters, final state, trajectory, topological charge) right here.
Each simulation rents a dedicated GPU (RTX 4090 / A100 class) on demand from vast.ai; the instance is created for your job and destroyed right after it (typical GPU cost: a few cents per job). The chat assistant is powered by the Anthropic Claude API (paid per token).
Both the GPU time and the Claude API tokens are funded and paid personally by Dieter Suess (University of Vienna) to keep this demo free and open for everyone. Please use it fairly — one job runs at a time, jobs are capped at 2.5 M cells / 20 ns. Note: on-demand GPU provisioning adds a few minutes before your job starts.
Your browser gets a random secret session ID (stored locally, like a cookie — no account, no personal data). Everything is scoped to it:
The session lives in this browser: a different device or cleared browser data starts a fresh, empty session. "New chat" clears the conversation but keeps your jobs.
Permalloy switching dynamics, 500×125×3 nm — the classic code benchmark.
Cortés-Ortuño standard problem: skyrmion relaxes to R ≈ 22 nm, Q = −1.
Your own box, material and field — see where the magnetisation settles.