NeuralMag

https://docs.neuralmag.org/
Admin
NeuralMag is a differentiable micromagnetic simulator built on JAX and PyTorch with a SymPy form...
Tokens:142,549
Snippets:835
Trust Score:3.4
Update:1 month ago
Show doc for...
Tokens:
Raw
### Complete simulation script

Source: https://docs.neuralmag.org/_sources/user_guide/getting_started.ipynb.txt

A full example script demonstrating the complete workflow from mesh setup to time integration.

```python
import neuralmag as nm

# setup mesh and state
mesh = nm.Mesh((100, 25, 1), (5e-9, 5e-9, 3e-9))
state = nm.State(mesh)

# setup material and m0
state.material.Ms = 8e5
state.material.A = 1.3e-11
state.material.alpha = 0.02

# initialize nodal vector functions for magneization and external field
state.m = nm.VectorFunction(state).fill((0.5**0.5, 0.5**0.5, 0))
h_ext = nm.VectorFunction(state).fill([-19576.0, 3421.0, 0.0], expand=True)

# register effective field contributions
nm.ExchangeField().register(state, "exchange")
nm.DemagField().register(state, "demag")
nm.ExternalField(h_ext).register(state, "external")
nm.TotalField("exchange", "demag").register(state)

# relax to s-state
llg = nm.LLGSolver(state)
llg.relax()

state.write_vti("m", "sstate.vti")

# add external field to perform switch
nm.TotalField("exchange", "demag", "external").register(state)
llg.reset()

logger = nm.Logger("data", ["t", "m"], ["m"])
while state.t < 1e-9:
    logger.log(state)
    llg.step(1e-11)
```

--------------------------------

### Initialize Mesh objects

Source: https://docs.neuralmag.org/reference/state.html

Examples of creating 3D and 2D meshes with specific cell counts and sizes.

```python
>>> # 3D mesh with one cell thickness, leading to 2 nodes in z-direction
>>> mesh_3d = Mesh((100, 25, 1), (5e-9, 5e-9, 3e-9))
>>> # 2D mesh with one cell thickness, leading to 1 node in z-direction
>>> mesh_2d = Mesh((100, 25), (5e-9, 5e-9, 3e-9))
```

--------------------------------

### Setup mesh and state

Source: https://docs.neuralmag.org/examples/standard-problem-4.html

Creates a mesh with 100x25x1 cells of size 5x5x3 nm and initializes a simulation state.

```python
mesh = nm.Mesh([100, 25, 1], [5e-9, 5e-9, 3e-9], [0.0, 0.0, 0.0])
state = nm.State(mesh)
```

--------------------------------

### Setup mesh and state with material parameters

Source: https://docs.neuralmag.org/examples/time-dependent-optimization_torch.html

Creates a 2x2x2 mesh with 5 nm cell size, initializes the state, sets material parameters (Ms, A, Ku, Ku_axis, alpha), and sets the initial magnetization to the z-direction.

```python
mesh = nm.Mesh((2, 2, 2), (5e-9, 5e-9, 5e-9))
state = nm.State(mesh)

state.material.Ms = 8e5
state.material.A = 1.3e-11
state.material.Ku = 1e5
state.material.Ku_axis = [0, 0, 1]
state.material.alpha = 0.1

state.m = nm.VectorFunction(state).fill((0, 0, 1))

```

--------------------------------

### Install NeuralMag with JAX backend

Source: https://docs.neuralmag.org/index.html

Use this command to install the package with the JAX backend, which requires Python 3.10 or higher.

```bash
pip install "neuralmag[jax]"
```

--------------------------------

### Install NeuralMag with PyTorch backend

Source: https://docs.neuralmag.org/index.html

Use this command to install the package with the PyTorch backend.

```bash
pip install "neuralmag[torch]"
```

--------------------------------

### Inverse-problem optimization loop in JAX

Source: https://docs.neuralmag.org/user_guide/dynamic_attributes.html

A complete example of an optimization loop using state.resolve in JAX.

```python
import jax, jax.numpy as jnp
import neuralmag as nm

state.rho_m = nm.CellFunction(state).fill(1.0)
state.rho   = nm.CellFunction(
    state, tensor=lambda rho_m: jnp.where(mask, rho_m, state.eps)
)

demag_func = state.resolve("h_demag", ["rho_m"])

def loss(rho_m):
    h = demag_func(rho_m ** 3)
    return -(h[10, 10, 12, 2] ** 2)

grad_loss = jax.grad(loss)

for step in range(N):
    g = grad_loss(state.rho_m.tensor)
    state.rho_m.tensor = jnp.clip(
        state.rho_m.tensor - lr * g, state.eps, 1.0
    )
```

--------------------------------

### Setup mesh and state

Source: https://docs.neuralmag.org/_sources/examples/time-dependent-optimization_torch.ipynb.txt

Defines the 3D mesh, material parameters, and sets the initial magnetization vector in the z-direction.

```python
mesh = nm.Mesh((2, 2, 2), (5e-9, 5e-9, 5e-9))
state = nm.State(mesh)

state.material.Ms = 8e5
state.material.A = 1.3e-11
state.material.Ku = 1e5
state.material.Ku_axis = [0, 0, 1]
state.material.alpha = 0.1

state.m = nm.VectorFunction(state).fill((0, 0, 1))
```

--------------------------------

### Initialize State and set magnetization

Source: https://docs.neuralmag.org/reference/state.html

Example of initializing a State object with a Mesh and setting magnetization using coordinate functions.

```python
state = nm.State(nm.Mesh((10, 10, 10), (1e-9, 1e-9, 1e-9)))
x, y, z = state.coordinates('nnn')

# initialize magnetization based on coordinate function
state.m = VectorFunction(state)
state.m.tensor[..., 0] = torch.sin(x/20e-9)
state.m.tensor[..., 1] = torch.cos(x/20e-9)
```

--------------------------------

### Setup: Import neuralmag and create 2D mesh

Source: https://docs.neuralmag.org/examples/LI.html

Imports neuralmag and creates a 20x20 2D mesh with unit cell size 1e-9. Sets the default dtype to float64 and initializes the state with magnetization pointing along z.

```python
import matplotlib.pyplot as plt
import numpy as np

import neuralmag as nm

nm.config.dtype = "float64"

mesh = nm.Mesh((20, 20), (1e-9, 1e-9, 1e-9))
state = nm.State(mesh)

state.m = nm.VectorFunction(state).fill((0, 0, 1))

```

--------------------------------

### Registering a Bulk DMI Field

Source: https://docs.neuralmag.org/examples/LI.html

This example demonstrates how to set up a mesh and state, then register a BulkDMIField alongside an exchange field.

```python
mesh = nm.Mesh((10, 10, 10), (1e-9, 1e-9, 1e-9))
state = nm.State(mesh)

state.m = nm.VectorFunction(state).fill((0, 0, 1))

state.material.Ms = 0.86e6
state.material.alpha = 1
state.material.A = 1.3e-11
nm.ExchangeField().register(state, "exchange")


state.material.Db = 15e-3
nm.BulkDMIField().register(state, "bulk")

nm.TotalField("exchange", "bulk").register(state)
```

```text
2025-09-22 13:03:39 NeuralMag:INFO [Mesh] 3D, 10 x 10 x 10 (size = 1e-09 x 1e-09 x 1e-09)
2025-09-22 13:03:39 NeuralMag:INFO [NeuralMag] Set default device to 'TFRT_CPU_0'.
2025-09-22 13:03:39 NeuralMag:INFO [State] Running on device: TFRT_CPU_0 (dtype = float64, backend = jax)
2025-09-22 13:03:39 NeuralMag:INFO [ExchangeField] Register state methods (field: 'h_exchange', energy: 'E_exchange')
2025-09-22 13:03:39 NeuralMag:INFO [BulkDMIField] Register state methods (field: 'h_bulk', energy: 'E_bulk')
2025-09-22 13:03:39 NeuralMag:INFO [TotalField] Register state methods (field: 'h', energy: 'E')
```

--------------------------------

### LLGSolverJAX Relaxation Logs

Source: https://docs.neuralmag.org/_sources/examples/skyrmion-disk.ipynb.txt

Example output logs from the LLGSolverJAX showing the relaxation steps and final energy calculation.

```text
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 4.94813e+10) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 3.11297e+10) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 9.68685e+09) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 2.76109e+09) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 2.40585e+09) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 2.08952e+09) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 1.81648e+09) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 1.57754e+09) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 1.36861e+09) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 1.189e+09) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation step (max dm/dt = 1.03094e+09) 1/s[0m
2025-09-19 14:54:10 NeuralMag:INFO [1;37;34m[LLGSolverJAX] Relaxation finished, final energy E = -1.47908e-18 J[0m

```

--------------------------------

### Setup demagnetizing field

Source: https://docs.neuralmag.org/_sources/examples/topology-optimization_jax.ipynb.txt

Registers the demag-field class with the state and resolves a function for computing the field based on the material density.

```python
demag = nm.DemagField().register(state, "demag")
demag_func = state.resolve("h_demag", ["rho_m"])
```

--------------------------------

### LevelSetParameterization usage example

Source: https://docs.neuralmag.org/reference/parameterizations.html

Creates a level-set parameterization on a 2D mesh, registers it with the state, and resolves the density as a pure function of the design variables. The example shows the required imports and the typical setup for using LevelSetParameterization.

```python
>>> import neuralmag as nm
>>> state = nm.State(nm.Mesh((60, 40), (5e-9, 5e-9, 3e-9)))
>>> lsf = nm.LevelSetParameterization(12, 8, xlimits=(10, 50), ylimits=(5, 35))
>>> lsf.register(state)
>>> # pure, differentiable function of the design
>>> rho_func = state.resolve("rho", ["rbf_amplitudes"])
>>> rho = rho_func(state.rbf_amplitudes)
```

--------------------------------

### Initialize and step LLGSolver

Source: https://docs.neuralmag.org/reference/llg_solver.html

Demonstrates creating a state, registering an external Zeeman field, initializing the LLGSolver, and performing a single integration step.

```python
# create state with time and magnetization
state = nm.State(nm.Mesh((10, 10, 10), (1e-9, 1e-9, 1e-9)))
state.t = 0.0
state.m = nm.VectorFunction(state).fill((1, 0, 0))

# register constant Zeeman field as state.h
nm.ExternalField(torch.Tensor((0, 0, 8e5))).register(state, "")

# initialize LLGSolver
llg = LLGSolver(state)

# perform integration step
llg.step(1e-12)
```

--------------------------------

### state.resolve pre-binding

Source: https://docs.neuralmag.org/user_guide/dynamic_attributes.html

Example of pre-binding all arguments except the design variable.

```python
# Pre-binding everything except the design variable
state.a = 2.0
state.b = 4.0
c = lambda a, b: a * b
func = state.resolve(c, ["a"])   # b is bound to 4.0 from state
assert func(1.0) == 4.0
```

--------------------------------

### state.remap argument renaming

Source: https://docs.neuralmag.org/user_guide/dynamic_attributes.html

Example of using state.remap to rename function arguments.

```python
def f(a, b):
    return a + b

g = state.remap(f, {"a": "x", "b": "y"})
# g(x, y) is now identical to f(a, b)
```

--------------------------------

### Set up AdaBelief optimizer

Source: https://docs.neuralmag.org/_sources/examples/level-set-optimization_jax.ipynb.txt

Initializes the AdaBelief optimizer with a learning rate of 0.01 and defines the JIT-compiled optimization step function.

```python
optim = optax.adabelief(0.01)
opt_state = optim.init(state.rbf_amplitudes)


@eqx.filter_jit
def make_step(amps, opt_state):
    loss, grads = grad_loss(amps)
    updates, opt_state = optim.update(grads, opt_state)
    amps = eqx.apply_updates(amps, updates)
    return loss, amps, opt_state
```

--------------------------------

### state.resolve dependency injection

Source: https://docs.neuralmag.org/user_guide/dynamic_attributes.html

Example of swapping a sub-dependency on the fly using the inject parameter.

```python
# Swapping a sub-dependency on the fly
state.a = 2.0
state.b = 4.0
c = lambda a, b: a * b
func = state.resolve(c, ["e"], inject={"b": lambda e: 2 * e})
assert func(1.0) == 4.0          # c(a, b(e)) = 2 * (2*1) = 4
```

--------------------------------

### Setup mesh and state with material parameters

Source: https://docs.neuralmag.org/examples/level-set-optimization_torch.html

Creates a 2D mesh of 32x32 cells with 5 nm size, initializes a state, and sets material parameters: Ms=8e5, A=1.3e-11, alpha=1.0. The damping alpha=1 makes the dynamics strongly overdamped so the final state approximates equilibrium.

```python
mesh = nm.Mesh((32, 32), (5e-9, 5e-9, 5e-9))
state = nm.State(mesh)

state.material.Ms = 8e5
state.material.A = 1.3e-11
state.material.alpha = 1.0
```

--------------------------------

### Import libraries and configure backend

Source: https://docs.neuralmag.org/_sources/examples/time-dependent-optimization_torch.ipynb.txt

Initializes the environment by importing necessary libraries and setting the NeuralMag backend to 'torch' with reduced FEM quadrature order.

```python
import matplotlib.pyplot as plt
import numpy as np
import torch
from scipy import constants
from tqdm import tqdm

import neuralmag as nm

nm.config.backend = "torch"
nm.config.fem["n_gauss"] = 1
```

--------------------------------

### Initialize LLGSolver

Source: https://docs.neuralmag.org/_sources/examples/time-dependent-optimization_torch.ipynb.txt

Sets up the LLGSolver with the state and specifies 'angles' as parameters for gradient computation.

```python
llg = nm.LLGSolver(state, parameters=["angles"])
```

--------------------------------

### Field registration logs

Source: https://docs.neuralmag.org/examples/magnetization_reconstruction_from_nv_height_fit_jax.html

Example output logs generated after registering the physical field terms to the state.

```text
2026-09-07 08:36:03 NeuralMag:INFO [ExchangeField] Register state methods (field: 'h_exchange', energy: 'E_exchange', energy density: 'e_exchange')
2026-09-07 08:36:03 NeuralMag:INFO [DemagField] Register state methods (field: 'h_demag', energy: 'E_demag', energy density: 'e_demag')
2026-09-07 08:36:03 NeuralMag:INFO [DemagField]: building the demag tensor on the CPU (NumPy). Set NM_JAX_GPU_SETUP=1 or pass gpu_setup=True (requires PyTorch and a CUDA GPU) to build it on the GPU -- much faster for large meshes.
2026-09-07 08:36:03 NeuralMag:INFO [DemagField]: Set up demag tensor
2026-09-07 08:36:03 NeuralMag:INFO [InterfaceDMIField] Register state methods (field: 'h_idmi', energy: 'E_idmi', energy density: 'e_idmi')
2026-09-07 08:36:03 NeuralMag:INFO [UniaxialAnisotropyField] Register state methods (field: 'h_aniso', energy: 'E_aniso', energy density: 'e_aniso')
2026-09-07 08:36:03 NeuralMag:INFO [TotalField] Register state methods (field: 'h', energy: 'E', energy density: 'e')
```

--------------------------------

### Set up AdaBelief optimizer and define make_step with eqx.filter_jit

Source: https://docs.neuralmag.org/examples/level-set-optimization_jax.html

Sets up the AdaBelief optimizer with a learning rate of 0.01 and defines a JIT-compiled step function that computes the loss and gradients, updates the optimizer state, and applies the updates to the RBF amplitudes. The small learning rate is important because adaptive optimizers keep moving the design at roughly the learning rate even after convergence, and larger rates can drift the converged design onto configurations where the relaxed magnetization flips between degenerate ±y equilibria, causing spikes in the loss curve.

```python
optim = optax.adabelief(0.01)
opt_state = optim.init(state.rbf_amplitudes)


@eqx.filter_jit
def make_step(amps, opt_state):
    loss, grads = grad_loss(amps)
    updates, opt_state = optim.update(grads, opt_state)
    amps = eqx.apply_updates(amps, updates)
    return loss, amps, opt_state

```

--------------------------------

### Define ExchangeField energy density

Source: https://docs.neuralmag.org/user_guide/form_compiler.html

Example of implementing a FieldTerm subclass by defining an energy density function e_expr.

```python
from neuralmag.common.engine import N, Variable, dV

def e_expr(m, dim, options):
    A = Variable("material__A", "c" * dim)   # cell-constant exchange constant
    return A * (m.diff(N.x).dot(m.diff(N.x))
                + m.diff(N.y).dot(m.diff(N.y))
                + m.diff(N.z).dot(m.diff(N.z))) * dV(dim)
```

--------------------------------

### Setup mesh and state with material parameters (NeuralMag)

Source: https://docs.neuralmag.org/examples/level-set-optimization_jax.html

Creates a 32x32 mesh with 5 nm cell size and a state, then sets material parameters: saturation magnetization Ms=8e5 A/m, exchange constant A=1.3e-11 J/m, and damping alpha=1.0. The damping is strongly overdamped so the magnetization dynamics are not important; only the final state matters. The output shows the mesh creation and the state running on CPU with float64 precision and the JAX backend.

```python
mesh = nm.Mesh((32, 32), (5e-9, 5e-9, 5e-9))
state = nm.State(mesh)

state.material.Ms = 8e5
state.material.A = 1.3e-11
state.material.alpha = 1.0

```

```text
2026-09-03 15:09:19 NeuralMag:INFO [Mesh] 2D, 32 x 32 (size = 5e-09 x 5e-09 x 5e-09)
2026-09-03 15:09:20 NeuralMag:INFO [NeuralMag] Set default device to 'cpu:0'.
2026-09-03 15:09:20 NeuralMag:INFO [State] Running on device: cpu:0 (dtype = float64, backend = jax)

```

--------------------------------

### state.resolve dependency chain

Source: https://docs.neuralmag.org/user_guide/dynamic_attributes.html

Example of walking a dependency chain where c depends on b, and b depends on a.

```python
# Walking a chain: c depends on b, b depends on a
state.a = 1.0
state.b = lambda a: 2.0 * a
c = lambda b: 2.0 * b
func = state.resolve(c)          # signature: func(a)
assert func(1.0) == 4.0          # because b = 2*a, c = 2*b
```

--------------------------------

### Set up mesh, state and geometry

Source: https://docs.neuralmag.org/examples/synthetic-antiferromagnet.html

Creates a 10x10x3 mesh with 1 nm cell size and a state. Sets the middle layer as a nonmagnetic spacer by assigning the density field to machine epsilon (state.eps) to avoid division-by-zero errors.

```python
# setup mesh and state
mesh = nm.Mesh((10, 10, 3), (1e-9, 1e-9, 1e-9))
state = nm.State(mesh)

# set empty spacer layer
rho = np.ones(mesh.n)
rho[:, :, 1] = state.eps
state.rho = nm.CellFunction(state, tensor=state.tensor(rho))
```

--------------------------------

### Set up Adam optimizer

Source: https://docs.neuralmag.org/examples/level-set-optimization_torch.html

Initializes the Adam optimizer with a learning rate of 0.01. A small learning rate is recommended to prevent the design from drifting into unstable configurations.

```python
optimizer = torch.optim.Adam([amps], lr=0.01)
```
Tokens:
Raw