Introduction#

NeuralMag is a differentiable micromagnetic simulator built around a SymPy form compiler. It targets researchers who want to solve inverse problems — topology optimization, parameter identification, control problems — in addition to classical forward simulations. The code base sits on top of either PyTorch or JAX, so every quantity that the form compiler produces is automatically differentiable and runs on CPU or GPU without any code changes.

What you get#

  • Nodal finite-difference discretization on a regular cuboid grid (the default).

  • A cell-centred variant — FIC — that reuses exactly the same form compiler, handy whenever your design variables naturally live on cells (topology optimization, imaging-derived geometries, sharp material interfaces). See Discretization.

  • Symbolic energy functionals (exchange, demag, uni-/cubic anisotropy, interface and bulk DMI, interlayer exchange, external field) that are compiled to vectorized backend code at runtime.

  • A differentiable Landau–Lifshitz–Gilbert integrator and a small set of loggers.

  • Fully functional interface: everything that depends on the simulation state — material parameters, applied fields, even the geometry mask — can be a plain Python lambda of other state attributes, which makes gradients through optimization loops trivial.

Mental model#

A typical NeuralMag script wires together four kinds of objects:

Mesh  ──▶  State  ──▶  FieldTerm.register(state)   ──▶  LLGSolver / Loggers
                             │
                             ▼
                        state.h_demag, state.h_exchange, …
                        state.E_demag, state.E_exchange, …

  • Mesh — a regular cuboid grid: number of cells n, spacing dx, origin.

  • State — the central container. It owns the mesh, the device and dtype, material parameters (state.material.Ms, …), the magnetization state.m, and — via dynamic attributes — everything derived from them.

  • FieldTerm subclasses (ExchangeField, DemagField, …) each register two new dynamic attributes on the state: state.h_<name> for the effective field and state.E_<name> for the energy. You never touch the generated code yourself unless you write a new field term.

  • LLGSolver / Loggers — consume the registered state.h_* attributes to integrate the LLG equation and record data.

If you come from a classical FD code, the biggest surprise is that state.material.Ms does not have to be a tensor: it can be a lambda that depends on another state attribute. NeuralMag inspects the lambda’s argument names, walks the dependency graph lazily, and assembles the full computation only when you actually read state.h_demag. This is what makes inverse-problem loops pleasant — see Dynamic Attributes, resolve and remap.

Two discretization modes, one form compiler#

The default scheme places the magnetization on the nodes of the grid (trilinear Lagrange basis) while material parameters live on cells (piecewise constant). The FIC variant keeps the same nodal finite-difference field computation but stores m on cells; a pair of mass-lumped projections moves data between nodes and cells, generated by the same form compiler that handles every other term. Both schemes are covered in Discretization.

Backends and devices#

NeuralMag picks its backend and device from environment variables at import time:

  • NM_BACKEND — torch (default) or jax.

  • NM_DEVICE — cpu or cuda.

  • NM_DTYPE — float32 or float64.

For torch you can also override device / dtype per State via the device= and dtype= arguments. Switching backends at runtime is supported as long as the new backend is installed — every FieldTerm re-compiles transparently on first use, with generated code cached under ~/.cache/neuralmag.

Who this guide is for#

We assume you are comfortable with Python / NumPy and have seen micromagnetic simulations before. You do not need to know the internals of the form compiler to use NeuralMag — reading Introduction, Getting Started and Dynamic Attributes, resolve and remap is enough for most users. For background on the physics and the nodal finite-difference discretization, see Abert (2019) and Abert et al. (npj Comput. Mater. 11, 193, 2025).