Micromagnetics in JAX
Simulate magnetic textures — and differentiate through them.
Finite-difference and finite-element solvers behind one interface, FFT and fast-multipole stray fields, float64 on GPU, and automatic differentiation straight through the Landau–Lifshitz–Gilbert equation. Written in Python, fast enough for 216 million cells.
Two backends, one API
Swap jaxmag.fd for jaxmag.fe and the same driver script runs on a
regular grid or a tetrahedral mesh — FFT demag here, fast multipole there.
Differentiable physics
Field terms and whole trajectories are differentiable, so inverse design and parameter fitting become gradient descent instead of guesswork.
Geometry that follows a surface
Per-element surface normals drive anisotropy and interfacial DMI, so a measured AFM topography becomes a magnetic film with curvature effects included.
Scales past one card
Domain-decomposed FFT demag across GPUs, agreeing with the single-GPU result to under 1 A/m; 600³ cells on eight GPUs.
02 · Run online
Describe a simulation, get a report
Write what you want to simulate, in English or German. An assistant maps it onto a jaxmag template, rents a GPU, runs it, and returns a PDF report with the final state, the trajectory and the topological charge.
Permalloy switching
500 × 125 × 3 nm, field 1 — the reference benchmark every code is measured against.
Skyrmion relaxation
Cortés-Ortuño standard problem: settles at R ≈ 22 nm with Q = −1.
Your own box
Pick the size, material and field, and see where the magnetisation settles.
Your jobs
Only this browser sees your chat, jobs and reports.
03 · Meshing service
Turn your geometry into a mesh
Upload a measured topography or an existing mesh and get back a tetrahedral mesh
for jaxmag.fe, built with gmsh — with a preview, the element count and the
curvature-induced DMI to expect from it.
Then use it locally
# the download contains topo_verts.npy, topo_tets.npy and mesh.msh import numpy as np from jaxmag.fe import FEMesh, fix_orientations v = np.load('topo_verts.npy'); t = np.load('topo_tets.npy').astype('int32') mesh = FEMesh(v, fix_orientations(v, t), mu0_Ms=4e-7*np.pi*860e3, scale_len=1e-9)
04 · Viewer
Watch the magnetisation move
Every simulation keeps 24 snapshots of the full magnetisation. Play them back, scrub through time, switch between the components and read the topological charge as it changes.
05 · Files
Everything you have made here
Reports from your simulations, meshes from your uploads, and the jaxmag source bundle — in one place, scoped to this browser.
06 · Install & use
Get jaxmag running
Install from GitLab, or download the self-contained bundle with vendored dependencies, every example and the documentation.
# from GitLab (pulls jax-ode and jaxfmm automatically) pip install git+https://gitlab.com/dietersuess1/jaxmag.git # or from the downloaded bundle, offline 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
What the bundle documents
| file | what it covers |
|---|---|
| README.md | package structure and the unified FD/FE API, choosing an integrator |
| INSTALL.md | setup online and offline, GPU requirements, how to verify it works |
| METHODS.md | the physics and the numerics: field terms, demag splitting, BDF + GMRES, curved films |
| EXAMPLES.md | every example script, plus the skyrmion quickstart |
| TESTING.md | the test suite and what each test proves |
A skyrmion in ten lines
from jaxmag.fd import make_mesh, make_demag, make_exchange, make_dmi from jaxmag.core import make_llg_equation from jaxmag import MU0, GAMMA0 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, gamma0=GAMMA0), shape=mesh.shape)
07 · Benchmarks
Measured, on real hardware
Timings from the test suite and the scaling runs that ship in
examples/ — not projections.
| problem | backend | hardware | time |
|---|---|---|---|
| SP4, 5 nm cells, Dopri5 + demag extrapolation | FD | RTX 3060 | 0.32 s |
| SP4, 3 nm cells | FD | RTX 3060 | 0.68 s |
| SP4, 5 nm cells | FE | RTX 3060 | 2.1 s |
| SP4 relaxation, 14k nodes, BDF + gmres_approx | FE | A100 | 10.7 s |
| the same with exact Jacobian solves | FE | A100 | 130 s |
| 600³ = 216 M cells, 1 ns | FD | 8 GPUs | see repo |
08 · About
Who runs this, and who pays
jaxmag is developed in the Physics of Functional Materials group at the University of Vienna by Dieter Suess and contributors, and released as open source.
Rented per job
Simulations run on a GPU rented on demand from vast.ai and destroyed when the job ends. A typical job costs a few cents.
Claude API
The chat that turns your description into a simulation runs on the Anthropic Claude API, billed per token.
Paid personally
Both the GPU time and the API tokens are funded by Dieter Suess so this stays free and open for everyone. Please use it fairly.
Your session is yours
A random id in your browser scopes your chat, jobs, meshes and reports. Nobody else can list or download them; clearing browser data starts a fresh, empty session.