Reconstructing magnetization from NV magnetometry#

A scanning NV magnetometer measures a single number per pixel: the projection of the stray field onto the nitrogen-vacancy symmetry axis, \(H^\mathrm{dem}_\mathrm{NV} = \vec{H}^\mathrm{dem}(\vec{x}) \cdot \hat{n}_\mathrm{NV}\), recorded at an effective distance \(d_\mathrm{NV}\) above the sample. Recovering the magnetization \(\vec{m}\) from such a scan is a badly conditioned inverse problem: propagating a field over the distance \(d_\mathrm{NV}\) damps each spatial frequency \(k\) by \(e^{-k d_\mathrm{NV}}\), and the projection mixes the field components.

In this example we simulate a scan of a chiral labyrinth, then reconstruct the magnetization by minimizing \(J(\vec{m}, d_\mathrm{NV}) = \mathcal{L}_\mathrm{data}^2 + \lambda \, \mathcal{L}_\mathrm{energy}\), the relative error to the measurement plus the micromagnetic energy of the candidate state, with the forward model of Setescak et al., arXiv:2602.17180. The effective distance \(d_\mathrm{NV}\) is fitted alongside the magnetization, because an experiment never knows it precisely.

We assemble the reconstruction in four steps:

  1. Forward model. On a State holding mesh and material, we build forward(m, d_nv): it computes the stray field of a candidate magnetization, continues it upward to the distance \(d_\mathrm{NV}\), and projects it onto the NV axis. The output is the scan the instrument would record for that candidate.

  2. Regularization. The prior of this reconstruction is the micromagnetic energy of the candidate state. state.resolve("E", ["m"]) returns the energy of the registered field terms as a plain function E(m), so evaluating the prior is one function call inside the loss.

  3. Loss. loss(params, lam) collects the unknowns, the magnetization angles and the distance, into one argument and returns a single scalar: the relative data misfit plus \(\lambda\) times the energy.

  4. Optimization. We pass the loss to optax L-BFGS and iterate its update step. Automatic differentiation carries gradients from the scalar loss back through energy, propagation, and forward model, so no derivative is written by hand.

Import libraries#

XLA’s multi-threaded Eigen backend deadlocks jax 0.11 on the FFT convolutions used below when only two CPU cores are available (for example on CI runners), so we disable it before importing jax. The flag has no effect on GPU.

[1]:

import os

os.environ["XLA_FLAGS"] = os.environ.get("XLA_FLAGS", "") + " --xla_cpu_multi_thread_eigen=false"

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import optax
from scipy import constants

import neuralmag as nm

nm.config.backend = "jax"
nm.config.dtype = "float32"

/home/setescak/data3/my_envs/neuralmag_dev_venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
2026-09-07 08:36:01 NeuralMag:INFO [NeuralMag] Version 1.0.0
2026-09-07 08:36:02 NeuralMag:INFO [NeuralMag] Backend set to 'jax'.
2026-09-07 08:36:02 NeuralMag:INFO [NeuralMag] Set default dtype to 'float32'.

Create mesh and state#

The mesh is two cell layers: the magnetic film at the bottom, and one vacuum cell above it, declared through state.domains. Under the cell-averaged discretization the demag convolution returns the stray field averaged over each cell volume, so the vacuum cell holds the average of the stray field over the 2.4 nm directly above the film. The forward model below converts that average into the field at the sample surface and continues it upward to the measurement distance, which decouples the vertical discretization from \(d_\mathrm{NV}\).

The film is a single magnetic layer of 2.4 nm. The in-plane cell size of 3 nm resolves the exchange length \(l_\text{ex} = \sqrt{2A/\mu_0 M_s^2} = 3.5\,\text{nm}\). At the effective distance of 40 nm the texture’s dominant mode is damped by \(e^{-2\pi d_\mathrm{NV} / (63\,\text{nm})} \approx 0.02\), so only a heavily smoothed shadow of the texture reaches the scan.

[2]:

N_XY = 96  # in-plane cells
DX = 3e-9  # in-plane cell size [m]
DZ = 2.4e-9  # cell height [m]: film thickness, and vacuum-cell height
NZ_MATERIAL = 1  # magnetic cell layers
NZ = NZ_MATERIAL + 1  # + 1 vacuum cell on top
D_NV_TRUE = 40e-9  # effective distance of the simulated scan [m]

mesh = nm.Mesh((N_XY, N_XY, NZ), (DX, DX, DZ))
state = nm.State(mesh)

domains = np.ones(mesh.n, dtype=np.int32)
domains[:, :, NZ_MATERIAL:] = 0  # top cell layer is vacuum
state.domains.tensor = state.tensor(domains)

2026-09-07 08:36:02 NeuralMag:INFO [Mesh] 3D, 96 x 96 x 2 (size = 3e-09 x 3e-09 x 2.4e-09)
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
2026-09-07 08:36:02 NeuralMag:INFO [NeuralMag] Set default device to 'cpu:0'.
2026-09-07 08:36:02 NeuralMag:INFO [State] Running on device: cpu:0 (dtype = float32, backend = jax)

Set material parameters#

At zero field this material relaxes into a dense chiral labyrinth with both out-of-plane and in-plane character. The sign of \(D\) sets the Néel chirality of the walls, and the two chiralities are not equivalent to the instrument: one radiates its stray field into the half-space above the film, the other closes its flux there and produces an order of magnitude weaker scan. We use the radiating one.

[3]:

state.material.Ms = 1.4e6  # saturation magnetization [A/m]
state.material.A = 1.5e-11  # exchange constant [J/m]
state.material.Di = -3e-3  # interface DMI constant [J/m^2], sign sets the Neel chirality
state.material.Di_axis = [0, 0, 1]
state.material.Ku = 1.1e6  # uniaxial anisotropy [J/m^3]
state.material.Ku_axis = [0, 0, 1]

Ground truth#

The ground truth is precomputed and shipped with this example. The script `generate_ground_truth.py <magnetization_reconstruction_from_nv_height_fit/generate_ground_truth.py>`__ draws a random magnetization, smooths it over a few cells, and relaxes it to equilibrium with the overdamped LLG at zero field. Rerunning the script regenerates the file. Whatever texture formed is one the energy supports, which is the property the reconstruction’s prior leans on.

The file stores the magnetization as cell data on the two-layer mesh, with the vacuum cell zeroed. Field terms are registered after state.m exists: each term inspects the magnetization’s function space to decide which discretization to compile, and reading the cell data selects the cell-averaged discretization used throughout.

[4]:

state.m = state.read_vti("magnetization_reconstruction_from_nv_height_fit/ground_truth.vti", "m")
m_true = state.m.tensor

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

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')

This is the state we are going to try to recover.

[5]:

def plot_m(m, label, axes, titles=False):
    """One row of panels: m_x, m_y, m_z. The row label sits on the left."""
    m = np.asarray(m)
    for ax, comp, name in zip(axes, range(3), (r"$m_x$", r"$m_y$", r"$m_z$")):
        im = ax.imshow(m[..., comp].T, cmap="RdBu_r", origin="lower", vmin=-1, vmax=1)
        if titles:
            ax.set_title(name)
        ax.set_xticks([])
        ax.set_yticks([])
    axes[0].set_ylabel(label)
    return im


fig, axes = plt.subplots(1, 3, figsize=(9.5, 3.2), constrained_layout=True)
im = plot_m(m_true[:, :, 0, :], "ground truth", axes, titles=True)
fig.colorbar(im, ax=axes, shrink=0.85, pad=0.01)
plt.show()

The forward model#

The forward model follows the paper. NeuralMag evaluates the demag field as a convolution: it stores the demag kernel state.N_demag in Fourier space on a zero-padded grid (twice the mesh size along each axis, for open boundaries), multiplies it with the transformed magnetization, and crops the result back to the mesh. Here we keep the full padded field instead of the crop: the upward-continuation FFT then runs at the padded k-space resolution, and the propagated stray field cannot wrap around the periodic FFT boundary. state.N_demag is the library’s internal kernel storage, exposed after DemagField().register(), not yet documented public API.

From the padded field we take the vacuum-cell layer, which holds the stray field averaged over the cell height \(\Delta_z\). In the vacuum above the sample the field satisfies the Laplace equation, so each in-plane Fourier mode \(k\) decays as \(e^{-kz}\), and two exact factors carry the average to the measurement:

\[ \tilde{\vec{H}}^\mathrm{dem}(d_\mathrm{NV}) = \langle \tilde{\vec{H}}^\mathrm{dem} \rangle_{[z_0,\, z_0+\Delta_z]} \cdot \left( \frac{k\Delta_z}{1 - e^{-k\Delta_z}} \right) e^{-k \, d_\mathrm{NV}}\]

The first factor de-averages: it converts the cell average into the field at the sample surface \(z_0\). The second factor continues the surface field upward by \(d_\mathrm{NV}\). After the inverse FFT the field is cropped to the mesh and projected onto the NV axis. The distance enters only through the continuation factor, so forward keeps it as an argument and the fit below treats it as a free parameter.

[6]:

NV_THETA = np.deg2rad(54.74)  # NV axis, tilted in the xz plane
NV_AXIS = jnp.array([np.sin(NV_THETA), 0.0, np.cos(NV_THETA)])

S_PAD = (2 * N_XY, 2 * N_XY, 2 * NZ)  # zero-padded grid of the open-BC demag kernel
k_1d = 2.0 * jnp.pi * jnp.fft.fftfreq(2 * N_XY, d=DX)
K_PAD = jnp.sqrt(k_1d[:, None] ** 2 + k_1d[None, :] ** 2)

N_DEMAG = state.N_demag  # Fourier-space demag kernel, entries (2 N_XY, 2 N_XY, NZ + 1)
RHO_MS = state.rho.tensor * state.material.Ms  # Ms with the vacuum cell suppressed


def reconstruct_field_at_z0(h_fft, k, dz):
    """De-average the cell-averaged field to the sample surface at the cell's bottom face."""
    kd = k * dz
    scale = jnp.where(kd < 1e-3, 1.0 + 0.5 * kd + kd**2 / 12.0, kd / (1.0 - jnp.exp(-kd)))
    return h_fft * scale[:, :, None]


def propagate_in_free_space(h_fft, k, distance):
    return h_fft * jnp.exp(-k * distance)[:, :, None]


def h_dem(m, d_nv):
    """Demag field at the distance d_nv above the film surface, as a 3-vector map [A/m]."""
    hx = jnp.zeros(N_DEMAG[0][0].shape, dtype=jnp.complex64)
    hy = jnp.zeros_like(hx)
    hz = jnp.zeros_like(hx)
    for ax in range(3):
        m_fft = jnp.fft.rfftn(RHO_MS * m[..., ax], axes=(0, 1, 2), s=S_PAD)
        hx = hx + N_DEMAG[0][ax] * m_fft
        hy = hy + N_DEMAG[1][ax] * m_fft
        hz = hz + N_DEMAG[2][ax] * m_fft
    hx = jnp.fft.irfftn(hx, axes=(0, 1, 2), s=S_PAD)  # (2 N_XY, 2 N_XY, 2 NZ), no crop
    hy = jnp.fft.irfftn(hy, axes=(0, 1, 2), s=S_PAD)
    hz = jnp.fft.irfftn(hz, axes=(0, 1, 2), s=S_PAD)
    h_vac = jnp.stack([hx[:, :, NZ - 1], hy[:, :, NZ - 1], hz[:, :, NZ - 1]], axis=-1)
    h_k = jnp.fft.fft2(h_vac, axes=(0, 1))
    h_k = reconstruct_field_at_z0(h_k, K_PAD, DZ)
    h_k = propagate_in_free_space(h_k, K_PAD, d_nv)
    return jnp.fft.ifft2(h_k, axes=(0, 1)).real[:N_XY, :N_XY, :]  # crop AFTER propagating


def forward(m, d_nv):
    """The scan: the demag field at d_nv, projected onto the NV axis [A/m]."""
    return h_dem(m, d_nv) @ NV_AXIS

Synthetic measurement#

The measurement is the scan of the ground truth at \(d_\mathrm{NV} = 40\,\)nm, with 10 % Gaussian noise added relative to the standard deviation of the signal. The measurement grid coincides with the simulation grid, and simulating the data with the same operator we invert makes this an inverse crime, so the reconstruction below is correspondingly optimistic.

[7]:

rng = np.random.default_rng(0)
h_true = forward(m_true, D_NV_TRUE)
h_meas = jnp.asarray(h_true + 0.10 * float(jnp.std(h_true)) * rng.standard_normal(h_true.shape))

The scan is much smoother than the magnetization that produced it, and the wall structure has left little visible trace: that smoothing is the \(e^{-k d_\mathrm{NV}}\) factor, and it is the reason the inverse problem needs a prior.

[8]:

h_ka = np.asarray(h_meas) * 1e-3
h_vmax = float(np.abs(h_ka).max())

fig, ax = plt.subplots(figsize=(4.6, 4), constrained_layout=True)
im = ax.imshow(h_ka.T, cmap="RdBu_r", origin="lower", vmin=-h_vmax, vmax=h_vmax)
ax.set_xticks([])
ax.set_yticks([])
fig.colorbar(im, ax=ax, shrink=0.8, label=r"$H^\mathrm{dem}_\mathrm{NV}$ (kA/m)")
plt.show()

Parameterize the unknown#

The magnetization is a unit-vector field, so we optimize spherical angles: \(|\vec{m}| = 1\) then holds by construction. The film is a single cell layer, so one pair of angles per cell describes it, and the vacuum cell stays at zero. The reconstruction starts from a random configuration, and the clip keeps \(\theta\) away from the poles, where the gradient with respect to \(\phi\) vanishes.

The effective distance is the second unknown. It starts at \(d_{\mathrm{NV},0} = 25\,\)nm, well below the true 40 nm, and is carried as a single scalar in nanometres so its magnitude matches that of the angles. The loss reads its absolute value, so the optimizer can never propose a negative distance.

[9]:

def to_magnetization(params):
    sin_theta = jnp.sin(params["theta"])
    m_film = jnp.stack(
        [
            sin_theta * jnp.cos(params["phi"]),
            sin_theta * jnp.sin(params["phi"]),
            jnp.cos(params["theta"]),
        ],
        axis=-1,
    )
    return jnp.zeros((N_XY, N_XY, NZ, 3)).at[:, :, 0, :].set(m_film)


D_NV_INIT = 25e-9  # initial guess [m], deliberately off the true 40 nm

rng = np.random.default_rng(1)
n_cells = mesh.n[:2]
params_init = {
    "theta": jnp.asarray(np.arccos(np.clip(rng.uniform(-1.0, 1.0, n_cells), -0.995, 0.995))),
    "phi": jnp.asarray(rng.uniform(-np.pi, np.pi, n_cells)),
    "d_nv_nm": jnp.asarray(1e9 * D_NV_INIT),
}

The loss#

\[ J(\vec{m}, d_\mathrm{NV}) = \mathcal{L}_\mathrm{data}^2 + \lambda \, \mathcal{L}_\mathrm{energy}, \qquad \mathcal{L}_\mathrm{data} = \frac{\lVert H^\mathrm{dem}_\mathrm{NV}(\vec{m}, d_\mathrm{NV}) - H^\mathrm{meas} \rVert}{\lVert H^\mathrm{meas} \rVert}\]

\(\mathcal{L}_\mathrm{data}\) is the relative field error, and \(\mathcal{L}_\mathrm{energy}\) is the total micromagnetic energy as a density over the film, measured from the uniform out-of-plane state \(\vec{m} = \vec{e}_z\) and normalized by it. Both terms are dimensionless, so \(\lambda\) is dimensionless too. state.resolve("E", ["m"]) gives the energy of the registered TotalField as a differentiable function of the magnetization. The distance enters \(\mathcal{L}_\mathrm{data}\) only: the energy knows nothing about the instrument. The weight \(\lambda\) stays a plain argument of loss, which pays off in the optimizer below.

[10]:

E_total = state.resolve("E", ["m"])

CELL_VOLUME = DX * DX * DZ
V_FILM = N_XY * N_XY * NZ_MATERIAL * CELL_VOLUME

m_ref = jnp.zeros((N_XY, N_XY, NZ, 3)).at[:, :, 0, 2].set(1.0)  # uniform out-of-plane film
E_DENS_REF = float(E_total(m_ref)) / V_FILM  # signed reference density [J/m^3]
H_MEAS_NORM = float(jnp.linalg.norm(h_meas))


def L_data(m, d_nv):
    return jnp.linalg.norm(forward(m, d_nv) - h_meas) / H_MEAS_NORM


def L_energy(m):
    return (E_total(m) / V_FILM - E_DENS_REF) / abs(E_DENS_REF)


def loss(params, lam):
    m = to_magnetization(params)
    return L_data(m, 1e-9 * jnp.abs(params["d_nv_nm"])) ** 2 + lam * L_energy(m)

Reconstruct#

L-BFGS with a zoom line search. Nothing here is specific to magnetism: it is the standard optax quasi-Newton loop, with the whole physics hidden behind forward and E_total. Because \(\lambda\) is a traced argument of the jitted step, JAX compiles the step once and reuses it for every value of \(\lambda\) below.

[11]:

N_STEPS = 300

solver = optax.lbfgs(
    memory_size=20,
    linesearch=optax.scale_by_zoom_linesearch(max_linesearch_steps=10),
)


@jax.jit
def step(params, opt_state, lam):
    def loss_fn(p):
        return loss(p, lam)

    value, grad = optax.value_and_grad_from_state(loss_fn)(params, state=opt_state)
    updates, opt_state = solver.update(grad, opt_state, params, value=value, grad=grad, value_fn=loss_fn)
    return optax.apply_updates(params, updates), opt_state


def reconstruct(lam):
    params, opt_state = params_init, solver.init(params_init)
    for _ in range(N_STEPS):
        params, opt_state = step(params, opt_state, lam)
    return params

How much regularization?#

We reconstruct at three weights, starting each run from the same random initial guess: \(\lambda = 0\) with the prior off, a small \(\lambda = 10^{-3}\), and a near-optimal \(\lambda = 0.3\). For each we record the relative field error \(\mathcal{L}_\mathrm{data}\), the energy \(\mathcal{L}_\mathrm{energy}\), the converged distance \(d_\mathrm{NV}^\ast\), and, because this is a synthetic problem where the answer is known, the actual error of the reconstruction.

On real data the true error does not exist, so choosing \(\lambda\) takes an independent argument: the corner of an L-curve traced by a full sweep over \(\lambda\), a noise estimate, or cross-validation on held-out pixels. The optimum also moves with the sample, so it is not a constant to carry between experiments.

[12]:

LAM_SMALL = 1e-3
LAM_OPT = 0.3  # near-optimal weight, found by a sweep over lambda
LAMBDAS = [0.0, LAM_SMALL, LAM_OPT]

print(f"ground truth: L_data = {float(L_data(m_true, D_NV_TRUE)):.4f}  L_energy = {float(L_energy(m_true)):+.3f}")

results = {}
for lam in LAMBDAS:
    params = reconstruct(lam)
    m = to_magnetization(params)
    d_nv = 1e-9 * float(jnp.abs(params["d_nv_nm"]))
    h_fit = forward(m, d_nv)
    results[lam] = {
        "m": np.asarray(m[:, :, 0, :]),
        "h": np.asarray(h_fit),
        "d_nv_nm": 1e9 * d_nv,
        "l_data": float(jnp.linalg.norm(h_fit - h_meas) / H_MEAS_NORM),
        "l_energy": float(L_energy(m)),
        "error": float(jnp.sqrt(jnp.mean(jnp.sum((m[:, :, 0, :] - m_true[:, :, 0, :]) ** 2, -1)))),
    }
    print(
        f"lambda = {lam:<6g}  L_data = {results[lam]['l_data']:.4f}"
        f"  L_energy = {results[lam]['l_energy']:+.3f}"
        f"  d_NV = {results[lam]['d_nv_nm']:5.2f} nm"
        f"  RMS error in m = {results[lam]['error']:.3f}"
    )

ground truth: L_data = 0.0989  L_energy = -1.755
lambda = 0       L_data = 0.0945  L_energy = +8.276  d_NV = 24.88 nm  RMS error in m = 1.202
lambda = 0.001   L_data = 0.0978  L_energy = -1.647  d_NV = 39.83 nm  RMS error in m = 0.168
lambda = 0.3     L_data = 0.0984  L_energy = -1.755  d_NV = 40.00 nm  RMS error in m = 0.013

What the three runs show#

With the prior off (\(\lambda = 0\)) the fit reproduces the measurement and still fails: at this distance the data leave most degrees of freedom unconstrained, and the optimizer spends them fitting noise. The converged distance drifts to 25 nm, because lowering the distance sharpens the operator and the noise-fitting texture prefers it.

The regularized runs recover the true 40 nm from a start 15 nm below, to within 0.2 nm at both weights. At \(\lambda = 10^{-3}\) the labyrinth is already recovered at an RMS error of 0.17, overlaid with a short-wavelength ripple that the damping \(e^{-k d_\mathrm{NV}}\) hides from the scan. At \(\lambda = 0.3\) the ripple is gone and the RMS error drops to 0.01, at a data misfit indistinguishable from the weakly regularized run.

The last column shows what each reconstruction looks like to the instrument: the field its magnetization produces at its converged \(d_\mathrm{NV}^\ast\), with the measurement itself in the top row. The rows below it step through the regularization: none, small, near-optimal. The \(\lambda = 0\) row reproduces the measurement while its magnetization is speckled: the data alone does not pin down the texture completely, and the prior supplies the difference. At \(\lambda = 10^{-3}\) the labyrinth is already recovered, with the residual ripple visible in all three components. At \(\lambda = 0.3\) both the field and the magnetization match.

[13]:

rows = [("ground truth", np.asarray(m_true[:, :, 0, :]), np.asarray(h_meas), 1e9 * D_NV_TRUE)]
rows += [(rf"$\lambda$ = {lam:g}", results[lam]["m"], results[lam]["h"], results[lam]["d_nv_nm"]) for lam in LAMBDAS]

h_vmax = 1e-3 * max(np.abs(h).max() for _, _, h, _ in rows)
fig, axes = plt.subplots(len(rows), 4, figsize=(12.5, 3.1 * len(rows)), constrained_layout=True)
for i, (axrow, (label, m, h, d_nv_nm)) in enumerate(zip(axes, rows)):
    im = plot_m(m, label, axrow[:3], titles=(i == 0))
    im_h = axrow[3].imshow(1e-3 * h.T, cmap="RdBu_r", origin="lower", vmin=-h_vmax, vmax=h_vmax)
    axrow[3].set_xticks([])
    axrow[3].set_yticks([])
    axrow[3].set_xlabel(rf"$d_\mathrm{{NV}} = {d_nv_nm:.1f}$ nm")
    if i == 0:
        axrow[3].set_title(r"$H^\mathrm{dem}_\mathrm{NV}$ (kA/m)")
fig.colorbar(im, ax=axes[:, :3], shrink=0.5, pad=0.01)
fig.colorbar(im_h, ax=axes[:, 3:], shrink=0.5, pad=0.02)
plt.show()

Write the result#

write_vti stores the reconstruction for inspection in ParaView, as cell data alongside the material density that marks which cells are film and which are vacuum.

[14]:

state.m.tensor = jnp.zeros((N_XY, N_XY, NZ, 3)).at[:, :, 0, :].set(jnp.asarray(results[LAM_OPT]["m"]))
state.write_vti(["m", "rho"], "magnetization_reconstruction_from_nv_height_fit/m_reconstructed.vti")