<!-- neuralmag 1.0.0 (latest main); for a version-pinned copy run neuralmag.llm_context() -->

# Using NeuralMag — guide for AI assistants

You are helping someone write **NeuralMag** code. NeuralMag is a micromagnetic
simulation package for **inverse problems** in magnetism: nodal finite-difference
discretization on regular cuboid meshes, PyTorch **and** JAX backends, and a
SymPy form compiler that turns energy functionals into vectorized, branch-free,
**autodifferentiable** tensor kernels.

NeuralMag is niche; your training data may be thin or wrong about it. **Trust the
live "Installed package" facts prepended to this guide, and the actual symbols
from `help(nm.X)` / `dir(nm)`, over your prior assumptions.** If an API isn't
shown here or visible via introspection, do not invent it.

**Use the existing functionality; do not extend NeuralMag on your own.** Solve
the task with the shipped field terms, solvers, functions, and loggers. Do
**not** write new `FieldTerm`/solver subclasses, monkey-patch, or otherwise add
to the library unless the user explicitly asks. If you believe a needed piece is
genuinely missing, **stop and ask the user** rather than inventing it — the
existing API almost always covers the case.

---

## 1. The one idea that governs everything

NeuralMag compiles physics into **plain functions on tensors with no control
flow over runtime values** — no `if`/`for` over tensor data, mesh sizes, or
material values in anything the solver or the gradient touches. That straight-line
form is what makes `torch.compile` / `jax.jit` effective and what makes the whole
pipeline differentiable end-to-end (the reason inverse problems are tractable).

**What this means for the code you write:**

- Express spatial/material variation as **data** (tensors, per-domain values,
  `rho`), not as Python branches.
- Bind quantities that depend on other quantities as **lambdas of state
  attributes** (see section 4), not as imperative recomputation inside the loop.
- Anything that branches on a runtime tensor value breaks the trace and defeats
  both JIT and autodiff. Don't do it.

If you find yourself writing `if m[...] > 0:` or looping over cells, stop — there
is almost always a vectorized or per-domain way to express it.

---

## 2. Configuration (backend / device / dtype)

Selection is via **environment variables, read once at import**. The backend is
fixed for the session and **cannot be switched at runtime**:

| Env var            | Values                  | Meaning                                  |
|--------------------|-------------------------|------------------------------------------|
| `NM_BACKEND`       | `torch` / `jax`         | array backend (auto-detected if unset)   |
| `NM_DEVICE`        | `cpu` / `cuda`          | device (torch)                           |
| `NM_DTYPE`         | `float32` / `float64`   | precision                                |
| `NM_TORCH_COMPILE` | `True` / `False`        | apply `torch.compile` to kernels         |
| `NM_JAX_JIT`       | `True` / `False`        | apply `jax.jit` to kernels               |

```bash
NM_BACKEND=torch NM_DTYPE=float64 python my_sim.py
```

`dtype` may also be set in Python **before tensors are created**:

```python
import neuralmag as nm
nm.config.dtype = "float64"   # do this first; many problems need float64 to converge
```

Query the active config at runtime:

```python
nm.config.device         # current device
nm.config.dtype          # current dtype
nm.config.backend_name   # "torch" or "jax" (set once the backend is initialized)
```

> Pitfall: do **not** try to switch backends mid-script or mix torch and jax
> tensors. Pick one via `NM_BACKEND` and stay in it.

---

## 3. Canonical workflow (memorize this shape)

A complete simulation — muMAG **standard problem 4**, relax then switch:

```python
import neuralmag as nm
from scipy import constants

# 1. Mesh: n cells, cell size dx (m), optional origin. SP4 is a single-layer
#    thin film, so use a 2D mesh (100x25) under the default nodal discretization;
#    dx still has 3 entries and dx[2]=3nm sets the physical thickness.
mesh = nm.Mesh([100, 25], [5e-9, 5e-9, 3e-9], [0.0, 0.0, 0.0])

# 2. State holds mesh, material params, fields, and dynamic attributes.
state = nm.State(mesh)

# 3. Material parameters (SI units).
state.material.Ms = 8e5         # saturation magnetization [A/m]
state.material.A = 1.3e-11      # exchange constant [J/m]
state.material.alpha = 0.02     # Gilbert damping

# 4. Initial magnetization (a unit-vector field on the mesh).
state.m = nm.VectorFunction(state).fill((0.5**0.5, 0.5**0.5, 0.0))

# 5. Register effective-field contributions. Each .register(state, "name")
#    creates state.h_<name> and state.E_<name>. TotalField sums named terms
#    into state.h / state.E.
nm.ExchangeField().register(state, "exchange")
nm.DemagField().register(state, "demag")
nm.TotalField("exchange", "demag").register(state)

# 6. Solve. relax() drives the system to equilibrium by integrating the
#    damping-only (overdamped) LLG; step(dt) integrates the full LLG dynamics.
llg = nm.LLGSolver(state)
llg.relax()
state.write_vti(["m"], "sp4/s-state.vti")

# 7. Add an external field, rebuild the total field, reset the solver.
h_ext = nm.VectorFunction(state).fill(
    [-24.6e-3 / constants.mu_0, 4.3e-3 / constants.mu_0, 0.0]
)
nm.ExternalField(h_ext).register(state, "external")
nm.TotalField("exchange", "demag", "external").register(state)
llg.reset()

# 8. Time integration. state.t is the current time; advance it with step(dt).
logger = nm.Logger("sp4", ["t", "m"], ["m"])
while state.t < 1e-9:
    llg.step(1.0e-11)
    logger.log(state)
```

Key objects and methods (verify exact signatures with `help`):

- `nm.Mesh(n, dx, origin=...)` — `n` = cell counts (1D/2D/3D), `dx` = spacing in m
  (**always 3 entries**, even for 1D/2D meshes — the trailing ones set the
  thickness/width). For a **single-layer thin film** under the default nodal
  discretization, use a **2D mesh** `(nx, ny)` rather than a 3D `(nx, ny, 1)`:
  it has one node layer in z instead of two redundant ones, and `dx[2]` still
  carries the physical thickness (demag included). Use a full 3D mesh only when
  the magnetization actually varies through the thickness (multiple cell layers).
- `nm.State(mesh)` — central container.
- `state.material.<name> = ...` — material parameters (proxy onto the state).
- `nm.VectorFunction(state)` / `nm.Function(state, spaces)` — **nodal** (node-space)
  fields; this is the **default discretization**, use it for `state.m` and other
  fields unless cell-space is explicitly requested. `nm.CellFunction(state)` /
  `nm.VectorCellFunction(state)` are the cell-space (FIC) variants — material
  parameters (`Ms`, `A`, `Ku`, ...) and `rho` are cell fields, but `state.m`
  should be nodal `VectorFunction` by default. All take `.fill(value)` or `tensor=...`.
- `<FieldTerm>().register(state, name)` -> `state.h_<name>`, `state.E_<name>`.
- `nm.TotalField(*names).register(state)` -> `state.h`, `state.E`.
- `nm.LLGSolver(state)` -> `.relax(tol, dt)` (overdamped relaxation to equilibrium;
  `tol` is the `max|dm/dt|` stopping criterion in rad/s), `.step(dt)`, `.reset()`;
  `state.t` tracks the current time.
- `nm.Logger(dir, scalars, fields)` / `nm.ScalarLogger(path, columns)` -> `.log(state)`;
  `Logger` also has `.is_resumable()` / `.resume(state)` to continue an interrupted
  run from existing log files (see section 7.4).
- `state.write_vti(["m", ...], path)` — ParaView output.

The available field-term classes are listed in the live header above (e.g.
`ExchangeField`, `DemagField`, `UniaxialAnisotropyField`, `CubicAnisotropyField`,
`BulkDMIField`, `InterfaceDMIField`, `InterlayerExchangeField`, `ExternalField`,
`OerstedField`, `LIField`).

### Material parameters (SI units)

NeuralMag uses **strict SI units** throughout. Set these on `state.material.<name>`
(scalar, per-domain `fill_by_domain`, cell field, or lambda — see sections 4–5):

| Attribute | Meaning | SI unit |
|-----------|---------|---------|
| `Ms` | saturation magnetization | A/m |
| `A` | exchange constant | J/m |
| `Ku` (+ `Ku_axis`) | uniaxial anisotropy constant | J/m³ |
| `Kc` (+ `Kc_axis1/2/3`) | cubic anisotropy constant | J/m³ |
| `Db` | bulk DMI constant | J/m² |
| `Di` (+ `Di_axis`) | interface DMI constant | J/m² |
| `iA` | interlayer (RKKY) exchange constant | J/m² |
| `alpha` | Gilbert damping | dimensionless |

> **Field-unit convention.** Magnetic fields are **A/m**, not Tesla, and current
> density (`state.j`, for `OerstedField`) is **A/m²**. A field given in Tesla `B`
> is converted by dividing by `constants.mu_0`: e.g.
> `[-24.6e-3 / constants.mu_0, ...]` in section 3 and the time-dependent field in
> section 7.2. Likewise `Ms = 1.0 / constants.mu_0` corresponds to `μ₀·Ms = 1 T`.

---

## 4. Dynamic attributes (the lambda mechanism)

Almost any state attribute can be a **lambda whose parameter names are other
state attributes**. On read, NeuralMag resolves the lambda, fetches the named
attributes from the state, and calls it (lazily, memoizing the call graph):

```python
# temperature-dependent Ms, bound by argument name "T"
Ms0, Tc = 8e5, 400.0
state.T = 300.0
state.material.Ms = lambda T: Ms0 * (1 - T / Tc) ** 1.5
# reading state.material.Ms now evaluates the lambda; reassigning state.T
# invalidates the cache so the next read picks up the new value.
```

Rules and pitfalls (these bite often):

- **Argument names must match existing state attribute names.**
  `lambda temperature: ...` fails if you stored it as `state.T`.
- **Don't capture tensors by closure.** Pass dependencies through the parameter
  list so the dependency graph can see them — otherwise `resolve`/`remap` and
  autodiff miss them.
- **Plain tensor/scalar assignment overwrites a lambda** (and `int`/`float`/`list`
  are auto-converted to backend tensors). Reassigning any attribute clears the
  resolved-function cache, so avoid reassigning structural attributes inside hot
  loops.
- A lambda can also produce a `Function`'s data:
  `nm.VectorFunction(state, tensor=lambda T: state.tensor([0, 0, 1 - T/Tc]))`.

### Inverse problems: `state.resolve` / `state.remap`

For optimization you want a **pure function of just the design variables**, with
everything else pre-bound. That is `state.resolve`:

```python
# pure function of the design tensor rho_m; mesh, Ms, relaxed m, ... are baked in
demag_func = state.resolve("h_demag", ["rho_m"])
# now demag_func(rho_m_tensor) is differentiable with jax.grad / torch.autograd
```

- Call `resolve` **once, outside** the optimization loop (it has compile cost).
- Inside the loop, update only the **tensor** (`state.rho_m.tensor = ...`), never
  the lambda graph, so the resolved closure stays valid.
- `state.remap(func, {"a": "x"})` just renames arguments (smaller cousin of
  `resolve`).

---

## 5. Multiple regions / materials

Two distinct tools — pick by whether the geometry is fixed or is itself the
design variable.

### Material density `rho` (the universal mask)

`state.rho` is a `CellFunction` that is `1.0` in magnetic material and
`state.eps` (machine epsilon) outside. Every generated energy/field form is
multiplied by `rho`, so non-magnetic cells drop out **without any explicit
branch**. **Never set `rho` to exactly 0** — that produces `inf`/`nan`; use
`state.eps`.

### Discrete regions: `add_domain` + `fill_by_domain`

```python
x, y, z = state.coordinates()
L0 = 30e-9
state.add_domain(1, abs(x) <= L0)   # region 1
state.add_domain(2, abs(x) >  L0)   # region 2
# (domain 0 is the implicit background; the first add_domain zeros the default)

# per-domain material values, indexed [background_0, region_1, region_2, ...]
state.material.Ms = nm.CellFunction(state).fill_by_domain([0.0, 8e5, 8e5])
state.material.A  = nm.CellFunction(state).fill_by_domain([0.0, 1.3e-11, 1.3e-11])
```

Pitfalls:

- `fill_by_domain` works on **cell-space functions only** (`nm.CellFunction`);
  nodal functions raise `NotImplementedError`.
- Domain IDs are **non-negative integers**; `0` is reserved for background.
- Re-adding an ID overwrites earlier marks — with overlapping conditions, the
  last call wins.
- `state.coordinates()` returns physical coordinate arrays for building masks.

### Manual / per-layer values

You can also build a NumPy array and wrap it, instead of using domains:

```python
import numpy as np

Ku = np.zeros(mesh.n)
Ku[:, :, 2:] = 1e7                          # only the top layer
state.material.Ku = nm.CellFunction(state, tensor=state.tensor(Ku))
state.material.Ku_axis = [1, 0, 0]
```

### Topology optimization (geometry as design variable)

Make `rho` a lambda of a continuous design field `rho_m`, then `resolve` against
it (see section 4) — that is what makes the geometry differentiable. For a 2D
design region with minimum-feature-size control and a binarizing projection, use
the shipped `nm.LevelSetParameterization` instead of rolling your own (section
7.5).

---

## 6. Pitfalls checklist (things assistants get wrong)

- NO control flow on tensor values (`if m[i] > 0`, Python `for` over cells).
  YES vectorize, or express variation via domains / `rho` / tensors.
- NO switching `NM_BACKEND` at runtime or mixing torch & jax tensors.
  YES one backend per process, set via env var.
- NO `float32` for a problem that needs accuracy.
  YES `nm.config.dtype = "float64"` **before** creating tensors.
- NO setting `rho` (or a region's density) to exactly `0`.
  YES use `state.eps`.
- NO inventing `State`/solver methods, or writing new field-term/solver
  subclasses to "extend" NeuralMag on your own.
  YES use existing functionality; check `dir(nm)`, `help(nm.State)`,
  `help(nm.LLGSolver)`; if something seems missing, ask the user first.
- NO cell-space (`VectorCellFunction` / FIC) for `state.m` by default.
  YES nodal `nm.VectorFunction` unless cell-space is explicitly requested.
- NO a 3D mesh with a single cell layer (`(nx, ny, 1)`) for a thin film under
  nodal discretization.
  YES a 2D mesh `(nx, ny)`; `dx[2]` sets the thickness.
- NO mismatched lambda arg names, or capturing tensors by closure.
  YES name parameters after state attributes; pass dependencies explicitly.
- NO calling `state.resolve` inside the optimization loop, or reassigning the
  lambda each step.
  YES resolve once; mutate only `.tensor`.
- NO expecting a scalars-only logger to resume — `Logger.resume`/`is_resumable`
  need a field logger.
  YES include a field (e.g. `fields=["m"]`) if you want to resume (section 7.4).
- NO assuming the `demos/` examples ship with `pip install` — they don't. The
  full example gallery is on the docs site (below).

---

## 7. Worked patterns (distilled from the example gallery)

These build on the canonical shape in section 3; each shows one pattern you will
reach for often. All are distilled from the shipped examples — see the gallery
(section 8) for the full versions.

### 7.1 Geometry via domains + DMI (2D skyrmion)

Shows a 2D mesh, a circular magnetic region via `add_domain`, and interface DMI +
uniaxial anisotropy.

```python
import neuralmag as nm
from scipy import constants

# 2D mesh (one node layer in z). The origin centers the disk on (0, 0).
mesh = nm.Mesh((50, 50), (2e-9, 2e-9, 0.6e-9), (-50e-9, -50e-9, 0))
state = nm.State(mesh)

state.material.Ms = 1.0 / constants.mu_0
state.material.A = 1.6e-11
state.material.Di = 4e-3            # interface DMI constant [J/m^2]
state.material.Di_axis = [0, 0, 1]
state.material.Ku = 510e3           # uniaxial anisotropy constant [J/m^3]
state.material.Ku_axis = [0, 0, 1]
state.material.alpha = 0.1

# 2D mesh -> coordinates() returns (x, y). rho = 1 inside the disk, eps outside.
x, y = state.coordinates()
state.add_domain(1, x**2 + y**2 < 50e-9**2)

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

nm.ExchangeField().register(state, "exchange")
nm.DemagField().register(state, "demag")
nm.InterfaceDMIField().register(state, "dmi")
nm.UniaxialAnisotropyField().register(state, "aniso")
nm.TotalField("exchange", "demag", "dmi", "aniso").register(state)

# scale_t rescales internal time units; relax(tol) stops at max|dm/dt| < tol [rad/s].
llg = nm.LLGSolver(state, scale_t=1e-12)
llg.relax(1e9)
state.write_vti(["m", "rho", "e"], "skyrmion-disk/m.vti")
```

### 7.2 Multilayer hysteresis: time-varying field + custom observable

Shows per-layer material via NumPy arrays, interlayer (RKKY) exchange, a
**time-dependent external field as a `lambda t`**, and a **custom logged
observable as a dynamic attribute**.

```python
import numpy as np
from scipy import constants

import neuralmag as nm

nm.config.dtype = "float64"   # set before creating tensors; needed for convergence

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

# Non-magnetic spacer in the middle cell layer: density = eps (never exactly 0).
rho = np.ones(mesh.n)
rho[:, :, 1] = state.eps
state.rho = nm.CellFunction(state, tensor=state.tensor(rho))

state.material.Ms = 1.0 / constants.mu_0
state.material.A = 1.3e-11
state.material.alpha = 1.0
state.material.iA = nm.Function(state, "ccn").fill(-0.5e-3, expand=True)  # RKKY [J/m^2]

# Uniaxial anisotropy only in the top ferromagnetic layer (pins it).
Ku = np.zeros(mesh.n)
Ku[:, :, 2:] = 1e7
state.material.Ku = nm.CellFunction(state, tensor=state.tensor(Ku))
state.material.Ku_axis = [1, 0, 0]

# Initial magnetization: top layer reversed along x.
m = np.zeros((11, 11, 4, 3))
m[..., 0] = 1.0
m[:, :, 2:, 0] = -1.0
state.m = nm.VectorFunction(state, tensor=state.tensor(m))

nm.ExchangeField().register(state, "exchange")
nm.InterlayerExchangeField(1, 2).register(state, "rkky")  # couples node layers 1 and 2
nm.ExternalField(
    lambda t: t * state.tensor([0.0, 5.0 / constants.mu_0 / 5e-9, 0.0])
).register(state, "external")
nm.UniaxialAnisotropyField().register(state, "aniso")
nm.TotalField("exchange", "rkky", "external", "aniso").register(state)

# Custom observable: mean magnetization over the bottom layer (backend-agnostic op).
state.m_bottom = lambda m: nm.config.backend.mean(m[:, :, :2, :], axis=(0, 1, 2))

llg = nm.LLGSolver(state)
logger = nm.ScalarLogger("saf/log.dat", ["h_external", "m_bottom"])
while state.t < 5e-9:
    logger.log(state)
    llg.step(1e-10)
```

### 7.3 Inverse design: topology optimization with autodiff

The reason NeuralMag exists. `rho` is made a differentiable function of a design
field `rho_m`; `state.resolve` turns the demag field into a pure function of
`rho_m` that `jax.grad` can differentiate. **Resolve once, outside the loop;
mutate only `.tensor` inside it.**

```python
import jax
import jax.numpy as jnp

import neuralmag as nm  # run with NM_BACKEND=jax

mesh = nm.Mesh((20, 20, 12), (5e-9, 5e-9, 5e-9))
state = nm.State(mesh)
state.material.Ms = 1.0
state.m = nm.VectorFunction(state).fill((0, 0, 1))

# rho depends on the design field rho_m; the top layers are forced to air.
state.rho = nm.CellFunction(state, tensor=lambda rho_m: rho_m.at[:, :, 10:].set(state.eps))
state.rho_m = nm.CellFunction(state).fill(1.0)

nm.DemagField().register(state, "demag")
demag_func = state.resolve("h_demag", ["rho_m"])   # pure function of rho_m — resolve ONCE

def loss(rho_m):
    # maximize H_z^2 at one point above the design region
    return -(demag_func(rho_m**3)[10, 10, 12, 2] ** 2)

grad_loss = jax.grad(loss)

lr = 1e3
for _ in range(100):
    g = grad_loss(state.rho_m.tensor)
    state.rho_m.tensor = jnp.clip(state.rho_m.tensor - lr * g, state.eps, 1.0)

state.write_vti("rho", "topology-optimization/rho.vti")
```

### 7.4 Resuming a simulation from existing logs

To restart a long run that was interrupted — appending to the existing `log.dat`
and `fields.pvd` instead of recomputing — use `logger.resume(state)`. **Resume
requires that a field is logged** (the magnetization field is what the state is
restored from); a scalars-only `Logger`/`ScalarLogger` cannot resume.

`resume(state)` restores `state.m` and `state.t` from the last recorded field
step, truncates and aligns all log files to that step, and continues appending. If
no resumable data exists it warns and leaves the state untouched, so the loop
simply starts from the beginning.

```python
logger = nm.Logger("sp4", ["t", "m"], ["m"])  # must log a field (here "m") to resume
if logger.is_resumable():
    logger.resume(state)   # restores state.m and state.t, aligns existing logs
while state.t < 1e-9:
    llg.step(1.0e-11)
    logger.log(state)
```

### 7.5 Level-set parameterization of `rho`

Optimizing `rho` cell-by-cell (7.3) gives no control over feature size and
produces grey, hard-to-manufacture designs. `nm.LevelSetParameterization`
instead parameterizes a rectangular design region of a **2D** mesh by a grid of
radial basis functions: their amplitudes are the design variables, a smooth
maximum combines them into a level-set field, an optional Gaussian filter sets
the minimum feature size, the field is renormalized to unit slope (the
"Eikonal" step, on by default), and a sigmoid of sharpness `a` projects it onto
`rho`. Because of the renormalization `a` is the inverse interface width in
cells (the 10–90 % transition of `rho` spans ≈ 4/a cells) regardless of the
filter width or the design; without it the filter and the smooth maximum
flatten the field's slope by a design-dependent factor and the same `a` gives
different widths. Ramping `a` upward during the optimization ("continuation")
anneals the design from smooth to binary.

```python
import numpy as np

import neuralmag as nm  # works under both backends

nm.config.dtype = "float64"   # REQUIRED on the jax backend: under float32 the
                              # gradients through llg.solve have been observed
                              # NaN (forward pass looks fine); torch gradients
                              # are fine under float32

state = nm.State(nm.Mesh((512, 64), (20e-9, 20e-9, 100e-9)))

# geometry: a fixed mask (1 = layout, eps = permanent void) the design sits in
geometry = np.ones(state.mesh.n)
geometry[:, :7] = state.eps

lsf = nm.LevelSetParameterization(
    40, 20,                          # RBF grid
    xlimits=(206, 306),              # design region, in CELL INDICES not metres
    ylimits=(7, 57),
    geometry=geometry,
    gaussian_sigma=3, gaussian_radius=9,   # minimum feature size, in cells
    sigmoid_a=2.0,                   # initial interface width ~ 4/a = 2 cells; ramp this up
    # eikonal=True (default): a means inverse interface width in cells;
    # eikonal=False projects the raw level-set field instead
)
lsf.register(state)                  # sets state.rho, .rbf_amplitudes, .lsf_sigmoid_a

# pure function of the design variables (and the sharpness, so it can be ramped)
rho_func = state.resolve("rho", ["rbf_amplitudes", "lsf_sigmoid_a"])
```

Registered attributes: `state.rbf_amplitudes` (design variables, shape
`(num_rbfs_x, num_rbfs_y)`, zeros unless you set them first),
`state.lsf_sigmoid_a`, `state.lsf_geometry`, and `state.rho`. Keeping
`lsf_sigmoid_a` a state attribute is what lets you thread it through
`nm.LLGSolver(state, parameters=["rbf_amplitudes", "lsf_sigmoid_a"])` and change
sharpness without re-resolving. With the default `eikonal=True`, `a = 4`
means a one-cell-wide interface and `a = 200` an effectively binary design;
ramp `a` geometrically. `rho` is clipped to `[state.eps, 1]`, so the
never-exactly-zero rule of section 5 holds automatically. On the jax backend do
not run this under float32: the forward pass looks healthy but gradients through
`llg.solve` come back NaN and silently destroy the optimization — `register()`
warns if a jax state is float32. Torch gradients flow normally under float32.

---

## 8. Where to look next

- Full docs & example gallery: https://neuralmag.gitlab.io/neuralmag/
- This guide, version-matched to the install: `neuralmag.llm_context()`
- Introspection: `dir(nm)`, `help(nm.State)`, `help(nm.LLGSolver)`,
  `help(nm.ExchangeField)`.