{ "cells": [ { "cell_type": "markdown", "id": "1589439c", "metadata": {}, "source": [ "# Reconstructing magnetization from NV magnetometry\n", "\n", "A scanning NV magnetometer measures a single number per pixel: the projection of the\n", "stray field onto the nitrogen-vacancy symmetry axis, $H^\\mathrm{dem}_\\mathrm{NV} =\n", "\\vec{H}^\\mathrm{dem}(\\vec{x}) \\cdot \\hat{n}_\\mathrm{NV}$, recorded at an effective\n", "distance $d_\\mathrm{NV}$ above the sample. Recovering the magnetization $\\vec{m}$ from\n", "such a scan is a badly conditioned inverse problem: propagating a field over the\n", "distance $d_\\mathrm{NV}$ damps each spatial frequency $k$ by $e^{-k d_\\mathrm{NV}}$, and\n", "the projection mixes the field components.\n", "\n", "In this example we simulate a scan of a chiral labyrinth, then reconstruct the\n", "magnetization by minimizing $J(\\vec{m}, d_\\mathrm{NV}) = \\mathcal{L}_\\mathrm{data}^2 +\n", "\\lambda \\, \\mathcal{L}_\\mathrm{energy}$, the relative error to the measurement plus the\n", "micromagnetic energy of the candidate state, with the forward model of\n", "[Setescak et al., arXiv:2602.17180](https://arxiv.org/abs/2602.17180). The effective\n", "distance $d_\\mathrm{NV}$ is fitted alongside the magnetization, because an experiment\n", "never knows it precisely.\n", "\n", "We assemble the reconstruction in four steps:\n", "\n", "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.\n", "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.\n", "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.\n", "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." ] }, { "cell_type": "markdown", "id": "7bc88885", "metadata": {}, "source": [ "## Import libraries\n", "\n", "XLA's multi-threaded Eigen backend deadlocks jax 0.11 on the FFT convolutions used below\n", "when only two CPU cores are available (for example on CI runners), so we disable it before\n", "importing jax. The flag has no effect on GPU." ] }, { "cell_type": "code", "execution_count": 1, "id": "dd8f5e7f", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:35:59.432271Z", "iopub.status.busy": "2026-09-07T08:35:59.432108Z", "iopub.status.idle": "2026-09-07T08:36:02.276346Z", "shell.execute_reply": "2026-09-07T08:36:02.275305Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/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\n", " from .autonotebook import tqdm as notebook_tqdm\n", "2026-09-07 08:36:01 NeuralMag:INFO \u001b[1;37;32m[NeuralMag] Version 1.0.0\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:02 NeuralMag:INFO \u001b[1;37;32m[NeuralMag] Backend set to 'jax'.\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:02 NeuralMag:INFO \u001b[1;37;32m[NeuralMag] Set default dtype to 'float32'.\u001b[0m\n" ] } ], "source": [ "import os\n", "\n", "os.environ[\"XLA_FLAGS\"] = os.environ.get(\"XLA_FLAGS\", \"\") + \" --xla_cpu_multi_thread_eigen=false\"\n", "\n", "import jax\n", "import jax.numpy as jnp\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import optax\n", "from scipy import constants\n", "\n", "import neuralmag as nm\n", "\n", "nm.config.backend = \"jax\"\n", "nm.config.dtype = \"float32\"" ] }, { "cell_type": "markdown", "id": "6aa4c4fc", "metadata": {}, "source": [ "## Create mesh and state\n", "\n", "The mesh is two cell layers: the magnetic film at the bottom, and one vacuum cell above\n", "it, declared through `state.domains`. Under the cell-averaged discretization the demag\n", "convolution returns the stray field averaged over each cell volume, so the vacuum cell\n", "holds the average of the stray field over the 2.4 nm directly above the film. The\n", "forward model below converts that average into the field at the sample surface and\n", "continues it upward to the measurement distance, which decouples the vertical\n", "discretization from $d_\\mathrm{NV}$.\n", "\n", "The film is a single magnetic layer of 2.4 nm. The in-plane cell size of 3 nm resolves\n", "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\n", "texture's dominant mode is damped by $e^{-2\\pi d_\\mathrm{NV} / (63\\,\\text{nm})} \\approx\n", "0.02$, so only a heavily smoothed shadow of the texture reaches the scan." ] }, { "cell_type": "code", "execution_count": 2, "id": "9ebd052f", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:02.278268Z", "iopub.status.busy": "2026-09-07T08:36:02.278008Z", "iopub.status.idle": "2026-09-07T08:36:02.781334Z", "shell.execute_reply": "2026-09-07T08:36:02.780616Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:02 NeuralMag:INFO \u001b[1;37;32m[Mesh] 3D, 96 x 96 x 2 (size = 3e-09 x 3e-09 x 2.4e-09)\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:02 NeuralMag:INFO \u001b[1;37;32m[NeuralMag] Set default device to 'cpu:0'.\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:02 NeuralMag:INFO \u001b[1;37;32m[State] Running on device: cpu:0 (dtype = float32, backend = jax)\u001b[0m\n" ] } ], "source": [ "N_XY = 96 # in-plane cells\n", "DX = 3e-9 # in-plane cell size [m]\n", "DZ = 2.4e-9 # cell height [m]: film thickness, and vacuum-cell height\n", "NZ_MATERIAL = 1 # magnetic cell layers\n", "NZ = NZ_MATERIAL + 1 # + 1 vacuum cell on top\n", "D_NV_TRUE = 40e-9 # effective distance of the simulated scan [m]\n", "\n", "mesh = nm.Mesh((N_XY, N_XY, NZ), (DX, DX, DZ))\n", "state = nm.State(mesh)\n", "\n", "domains = np.ones(mesh.n, dtype=np.int32)\n", "domains[:, :, NZ_MATERIAL:] = 0 # top cell layer is vacuum\n", "state.domains.tensor = state.tensor(domains)" ] }, { "cell_type": "markdown", "id": "74b45daf", "metadata": {}, "source": [ "## Set material parameters\n", "\n", "At zero field this material relaxes into a dense chiral labyrinth with both out-of-plane\n", "and in-plane character. The sign of $D$ sets the Néel chirality of the walls, and the two\n", "chiralities are not equivalent to the instrument: one radiates its stray field into the\n", "half-space above the film, the other closes its flux there and produces an order of\n", "magnitude weaker scan. We use the radiating one." ] }, { "cell_type": "code", "execution_count": 3, "id": "7ab716c1", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:02.783542Z", "iopub.status.busy": "2026-09-07T08:36:02.783355Z", "iopub.status.idle": "2026-09-07T08:36:02.787381Z", "shell.execute_reply": "2026-09-07T08:36:02.786759Z" } }, "outputs": [], "source": [ "state.material.Ms = 1.4e6 # saturation magnetization [A/m]\n", "state.material.A = 1.5e-11 # exchange constant [J/m]\n", "state.material.Di = -3e-3 # interface DMI constant [J/m^2], sign sets the Neel chirality\n", "state.material.Di_axis = [0, 0, 1]\n", "state.material.Ku = 1.1e6 # uniaxial anisotropy [J/m^3]\n", "state.material.Ku_axis = [0, 0, 1]" ] }, { "cell_type": "markdown", "id": "61be5490", "metadata": {}, "source": [ "## Ground truth\n", "\n", "The ground truth is precomputed and shipped with this example. The script\n", "[`generate_ground_truth.py`](magnetization_reconstruction_from_nv_height_fit/generate_ground_truth.py)\n", "draws a random magnetization, smooths it over a few cells, and relaxes it to equilibrium\n", "with the overdamped LLG at zero field. Rerunning the script regenerates the file.\n", "Whatever texture formed is one the energy supports, which is the property the\n", "reconstruction's prior leans on.\n", "\n", "The file stores the magnetization as cell data on the two-layer mesh, with the vacuum\n", "cell zeroed. Field terms are registered *after* `state.m` exists: each term inspects the\n", "magnetization's function space to decide which discretization to compile, and reading\n", "the cell data selects the cell-averaged discretization used throughout." ] }, { "cell_type": "code", "execution_count": 4, "id": "f023a939", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:02.788970Z", "iopub.status.busy": "2026-09-07T08:36:02.788773Z", "iopub.status.idle": "2026-09-07T08:36:03.859626Z", "shell.execute_reply": "2026-09-07T08:36:03.858598Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:03 NeuralMag:INFO \u001b[1;37;32m[ExchangeField] Register state methods (field: 'h_exchange', energy: 'E_exchange', energy density: 'e_exchange')\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:03 NeuralMag:INFO \u001b[1;37;32m[DemagField] Register state methods (field: 'h_demag', energy: 'E_demag', energy density: 'e_demag')\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:03 NeuralMag:INFO \u001b[1;37;32m[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.\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:03 NeuralMag:INFO \u001b[1;37;32m[DemagField]: Set up demag tensor\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:03 NeuralMag:INFO \u001b[1;37;32m[InterfaceDMIField] Register state methods (field: 'h_idmi', energy: 'E_idmi', energy density: 'e_idmi')\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:03 NeuralMag:INFO \u001b[1;37;32m[UniaxialAnisotropyField] Register state methods (field: 'h_aniso', energy: 'E_aniso', energy density: 'e_aniso')\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2026-09-07 08:36:03 NeuralMag:INFO \u001b[1;37;32m[TotalField] Register state methods (field: 'h', energy: 'E', energy density: 'e')\u001b[0m\n" ] } ], "source": [ "state.m = state.read_vti(\"magnetization_reconstruction_from_nv_height_fit/ground_truth.vti\", \"m\")\n", "m_true = state.m.tensor\n", "\n", "nm.ExchangeField().register(state, \"exchange\")\n", "nm.DemagField().register(state, \"demag\")\n", "nm.InterfaceDMIField().register(state, \"idmi\")\n", "nm.UniaxialAnisotropyField().register(state, \"aniso\")\n", "nm.TotalField(\"exchange\", \"demag\", \"idmi\", \"aniso\").register(state)" ] }, { "cell_type": "markdown", "id": "2b85d842", "metadata": {}, "source": [ "This is the state we are going to try to recover." ] }, { "cell_type": "code", "execution_count": 5, "id": "4a25dfac", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:03.861540Z", "iopub.status.busy": "2026-09-07T08:36:03.861340Z", "iopub.status.idle": "2026-09-07T08:36:03.984470Z", "shell.execute_reply": "2026-09-07T08:36:03.983315Z" } }, "outputs": [], "source": [ "def plot_m(m, label, axes, titles=False):\n", " \"\"\"One row of panels: m_x, m_y, m_z. The row label sits on the left.\"\"\"\n", " m = np.asarray(m)\n", " for ax, comp, name in zip(axes, range(3), (r\"$m_x$\", r\"$m_y$\", r\"$m_z$\")):\n", " im = ax.imshow(m[..., comp].T, cmap=\"RdBu_r\", origin=\"lower\", vmin=-1, vmax=1)\n", " if titles:\n", " ax.set_title(name)\n", " ax.set_xticks([])\n", " ax.set_yticks([])\n", " axes[0].set_ylabel(label)\n", " return im\n", "\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(9.5, 3.2), constrained_layout=True)\n", "im = plot_m(m_true[:, :, 0, :], \"ground truth\", axes, titles=True)\n", "fig.colorbar(im, ax=axes, shrink=0.85, pad=0.01)\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "ef77d9f4", "metadata": {}, "source": [ "## The forward model\n", "\n", "The forward model follows the paper. NeuralMag evaluates the demag field as a\n", "convolution: it stores the demag kernel `state.N_demag` in Fourier space on a\n", "zero-padded grid (twice the mesh size along each axis, for open boundaries), multiplies\n", "it with the transformed magnetization, and crops the result back to the mesh. Here we\n", "keep the **full padded field** instead of the crop: the upward-continuation FFT then\n", "runs at the padded k-space resolution, and the propagated stray field cannot wrap around\n", "the periodic FFT boundary. `state.N_demag` is the library's internal kernel storage,\n", "exposed after `DemagField().register()`, not yet documented public API.\n", "\n", "From the padded field we take the vacuum-cell layer, which holds the stray field\n", "averaged over the cell height $\\Delta_z$. In the vacuum above the sample the field\n", "satisfies the Laplace equation, so each in-plane Fourier mode $k$ decays as $e^{-kz}$,\n", "and two exact factors carry the average to the measurement:\n", "\n", "$$ \\tilde{\\vec{H}}^\\mathrm{dem}(d_\\mathrm{NV}) = \\langle \\tilde{\\vec{H}}^\\mathrm{dem}\n", "\\rangle_{[z_0,\\, z_0+\\Delta_z]} \\cdot \\left( \\frac{k\\Delta_z}{1 - e^{-k\\Delta_z}}\n", "\\right) e^{-k \\, d_\\mathrm{NV}} $$\n", "\n", "The first factor de-averages: it converts the cell average into the field at the sample\n", "surface $z_0$. The second factor continues the surface\n", "field upward by $d_\\mathrm{NV}$. After the inverse FFT the field is cropped to the mesh\n", "and projected onto the NV axis. The distance enters only through the continuation\n", "factor, so `forward` keeps it as an argument and the fit below treats it as a free\n", "parameter." ] }, { "cell_type": "code", "execution_count": 6, "id": "382ef04a", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:03.986360Z", "iopub.status.busy": "2026-09-07T08:36:03.986161Z", "iopub.status.idle": "2026-09-07T08:36:04.453209Z", "shell.execute_reply": "2026-09-07T08:36:04.452033Z" } }, "outputs": [], "source": [ "NV_THETA = np.deg2rad(54.74) # NV axis, tilted in the xz plane\n", "NV_AXIS = jnp.array([np.sin(NV_THETA), 0.0, np.cos(NV_THETA)])\n", "\n", "S_PAD = (2 * N_XY, 2 * N_XY, 2 * NZ) # zero-padded grid of the open-BC demag kernel\n", "k_1d = 2.0 * jnp.pi * jnp.fft.fftfreq(2 * N_XY, d=DX)\n", "K_PAD = jnp.sqrt(k_1d[:, None] ** 2 + k_1d[None, :] ** 2)\n", "\n", "N_DEMAG = state.N_demag # Fourier-space demag kernel, entries (2 N_XY, 2 N_XY, NZ + 1)\n", "RHO_MS = state.rho.tensor * state.material.Ms # Ms with the vacuum cell suppressed\n", "\n", "\n", "def reconstruct_field_at_z0(h_fft, k, dz):\n", " \"\"\"De-average the cell-averaged field to the sample surface at the cell's bottom face.\"\"\"\n", " kd = k * dz\n", " scale = jnp.where(kd < 1e-3, 1.0 + 0.5 * kd + kd**2 / 12.0, kd / (1.0 - jnp.exp(-kd)))\n", " return h_fft * scale[:, :, None]\n", "\n", "\n", "def propagate_in_free_space(h_fft, k, distance):\n", " return h_fft * jnp.exp(-k * distance)[:, :, None]\n", "\n", "\n", "def h_dem(m, d_nv):\n", " \"\"\"Demag field at the distance d_nv above the film surface, as a 3-vector map [A/m].\"\"\"\n", " hx = jnp.zeros(N_DEMAG[0][0].shape, dtype=jnp.complex64)\n", " hy = jnp.zeros_like(hx)\n", " hz = jnp.zeros_like(hx)\n", " for ax in range(3):\n", " m_fft = jnp.fft.rfftn(RHO_MS * m[..., ax], axes=(0, 1, 2), s=S_PAD)\n", " hx = hx + N_DEMAG[0][ax] * m_fft\n", " hy = hy + N_DEMAG[1][ax] * m_fft\n", " hz = hz + N_DEMAG[2][ax] * m_fft\n", " hx = jnp.fft.irfftn(hx, axes=(0, 1, 2), s=S_PAD) # (2 N_XY, 2 N_XY, 2 NZ), no crop\n", " hy = jnp.fft.irfftn(hy, axes=(0, 1, 2), s=S_PAD)\n", " hz = jnp.fft.irfftn(hz, axes=(0, 1, 2), s=S_PAD)\n", " h_vac = jnp.stack([hx[:, :, NZ - 1], hy[:, :, NZ - 1], hz[:, :, NZ - 1]], axis=-1)\n", " h_k = jnp.fft.fft2(h_vac, axes=(0, 1))\n", " h_k = reconstruct_field_at_z0(h_k, K_PAD, DZ)\n", " h_k = propagate_in_free_space(h_k, K_PAD, d_nv)\n", " return jnp.fft.ifft2(h_k, axes=(0, 1)).real[:N_XY, :N_XY, :] # crop AFTER propagating\n", "\n", "\n", "def forward(m, d_nv):\n", " \"\"\"The scan: the demag field at d_nv, projected onto the NV axis [A/m].\"\"\"\n", " return h_dem(m, d_nv) @ NV_AXIS" ] }, { "cell_type": "markdown", "id": "066006db", "metadata": {}, "source": [ "## Synthetic measurement\n", "\n", "The measurement is the scan of the ground truth at $d_\\mathrm{NV} = 40\\,$nm, with 10 %\n", "Gaussian noise added relative to the standard deviation of the signal. The measurement\n", "grid coincides with the simulation grid, and simulating the data with the same operator\n", "we invert makes this an inverse crime, so the reconstruction below is correspondingly\n", "optimistic." ] }, { "cell_type": "code", "execution_count": 7, "id": "09294074", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:04.455113Z", "iopub.status.busy": "2026-09-07T08:36:04.454912Z", "iopub.status.idle": "2026-09-07T08:36:06.369097Z", "shell.execute_reply": "2026-09-07T08:36:06.367932Z" } }, "outputs": [], "source": [ "rng = np.random.default_rng(0)\n", "h_true = forward(m_true, D_NV_TRUE)\n", "h_meas = jnp.asarray(h_true + 0.10 * float(jnp.std(h_true)) * rng.standard_normal(h_true.shape))" ] }, { "cell_type": "markdown", "id": "5d5dc057", "metadata": {}, "source": [ "The scan is much smoother than the magnetization that produced it, and\n", "the wall structure has left little visible trace: that smoothing is the\n", "$e^{-k d_\\mathrm{NV}}$ factor, and it is the reason the inverse problem needs a prior." ] }, { "cell_type": "code", "execution_count": 8, "id": "fbe349bc", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:06.370989Z", "iopub.status.busy": "2026-09-07T08:36:06.370790Z", "iopub.status.idle": "2026-09-07T08:36:06.406996Z", "shell.execute_reply": "2026-09-07T08:36:06.406070Z" } }, "outputs": [], "source": [ "h_ka = np.asarray(h_meas) * 1e-3\n", "h_vmax = float(np.abs(h_ka).max())\n", "\n", "fig, ax = plt.subplots(figsize=(4.6, 4), constrained_layout=True)\n", "im = ax.imshow(h_ka.T, cmap=\"RdBu_r\", origin=\"lower\", vmin=-h_vmax, vmax=h_vmax)\n", "ax.set_xticks([])\n", "ax.set_yticks([])\n", "fig.colorbar(im, ax=ax, shrink=0.8, label=r\"$H^\\mathrm{dem}_\\mathrm{NV}$ (kA/m)\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "04953c41", "metadata": {}, "source": [ "## Parameterize the unknown\n", "\n", "The magnetization is a unit-vector field, so we optimize spherical angles: $|\\vec{m}| = 1$\n", "then holds by construction. The film is a single cell layer, so one pair of angles per\n", "cell describes it, and the vacuum cell stays at zero. The reconstruction starts from a\n", "random configuration, and the clip keeps $\\theta$ away from the poles, where the\n", "gradient with respect to $\\phi$ vanishes.\n", "\n", "The effective distance is the second unknown. It starts at $d_{\\mathrm{NV},0} =\n", "25\\,$nm, well below the true 40 nm, and is carried as a single scalar in nanometres so\n", "its magnitude matches that of the angles. The loss reads its absolute value, so the\n", "optimizer can never propose a negative distance." ] }, { "cell_type": "code", "execution_count": 9, "id": "8bac2224", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:06.409247Z", "iopub.status.busy": "2026-09-07T08:36:06.408574Z", "iopub.status.idle": "2026-09-07T08:36:06.462719Z", "shell.execute_reply": "2026-09-07T08:36:06.461579Z" } }, "outputs": [], "source": [ "def to_magnetization(params):\n", " sin_theta = jnp.sin(params[\"theta\"])\n", " m_film = jnp.stack(\n", " [\n", " sin_theta * jnp.cos(params[\"phi\"]),\n", " sin_theta * jnp.sin(params[\"phi\"]),\n", " jnp.cos(params[\"theta\"]),\n", " ],\n", " axis=-1,\n", " )\n", " return jnp.zeros((N_XY, N_XY, NZ, 3)).at[:, :, 0, :].set(m_film)\n", "\n", "\n", "D_NV_INIT = 25e-9 # initial guess [m], deliberately off the true 40 nm\n", "\n", "rng = np.random.default_rng(1)\n", "n_cells = mesh.n[:2]\n", "params_init = {\n", " \"theta\": jnp.asarray(np.arccos(np.clip(rng.uniform(-1.0, 1.0, n_cells), -0.995, 0.995))),\n", " \"phi\": jnp.asarray(rng.uniform(-np.pi, np.pi, n_cells)),\n", " \"d_nv_nm\": jnp.asarray(1e9 * D_NV_INIT),\n", "}" ] }, { "cell_type": "markdown", "id": "7f84670a", "metadata": {}, "source": [ "## The loss\n", "\n", "$$ J(\\vec{m}, d_\\mathrm{NV}) = \\mathcal{L}_\\mathrm{data}^2 + \\lambda \\,\n", "\\mathcal{L}_\\mathrm{energy}, \\qquad \\mathcal{L}_\\mathrm{data} = \\frac{\\lVert\n", "H^\\mathrm{dem}_\\mathrm{NV}(\\vec{m}, d_\\mathrm{NV}) - H^\\mathrm{meas} \\rVert}{\\lVert\n", "H^\\mathrm{meas} \\rVert} $$\n", "\n", "$\\mathcal{L}_\\mathrm{data}$ is the relative field error, and\n", "$\\mathcal{L}_\\mathrm{energy}$ is the total micromagnetic energy as a density over the\n", "film, measured from the uniform out-of-plane state $\\vec{m} = \\vec{e}_z$ and normalized\n", "by it. Both terms are dimensionless, so $\\lambda$ is dimensionless too.\n", "`state.resolve(\"E\", [\"m\"])` gives the energy of the registered `TotalField` as a\n", "differentiable function of the magnetization. The distance enters\n", "$\\mathcal{L}_\\mathrm{data}$ only: the energy knows nothing about the instrument. The\n", "weight $\\lambda$ stays a plain argument of `loss`, which pays off in the optimizer\n", "below." ] }, { "cell_type": "code", "execution_count": 10, "id": "ede292f1", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:06.464630Z", "iopub.status.busy": "2026-09-07T08:36:06.464434Z", "iopub.status.idle": "2026-09-07T08:36:10.777523Z", "shell.execute_reply": "2026-09-07T08:36:10.776069Z" } }, "outputs": [], "source": [ "E_total = state.resolve(\"E\", [\"m\"])\n", "\n", "CELL_VOLUME = DX * DX * DZ\n", "V_FILM = N_XY * N_XY * NZ_MATERIAL * CELL_VOLUME\n", "\n", "m_ref = jnp.zeros((N_XY, N_XY, NZ, 3)).at[:, :, 0, 2].set(1.0) # uniform out-of-plane film\n", "E_DENS_REF = float(E_total(m_ref)) / V_FILM # signed reference density [J/m^3]\n", "H_MEAS_NORM = float(jnp.linalg.norm(h_meas))\n", "\n", "\n", "def L_data(m, d_nv):\n", " return jnp.linalg.norm(forward(m, d_nv) - h_meas) / H_MEAS_NORM\n", "\n", "\n", "def L_energy(m):\n", " return (E_total(m) / V_FILM - E_DENS_REF) / abs(E_DENS_REF)\n", "\n", "\n", "def loss(params, lam):\n", " m = to_magnetization(params)\n", " return L_data(m, 1e-9 * jnp.abs(params[\"d_nv_nm\"])) ** 2 + lam * L_energy(m)" ] }, { "cell_type": "markdown", "id": "8f6d30da", "metadata": {}, "source": [ "## Reconstruct\n", "\n", "L-BFGS with a zoom line search. Nothing here is specific to magnetism: it is the standard\n", "`optax` quasi-Newton loop, with the whole physics hidden behind `forward` and `E_total`.\n", "Because $\\lambda$ is a traced argument of the jitted `step`, JAX compiles the step once\n", "and reuses it for every value of $\\lambda$ below." ] }, { "cell_type": "code", "execution_count": 11, "id": "fb54ec23", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:10.779741Z", "iopub.status.busy": "2026-09-07T08:36:10.779527Z", "iopub.status.idle": "2026-09-07T08:36:10.785954Z", "shell.execute_reply": "2026-09-07T08:36:10.785091Z" } }, "outputs": [], "source": [ "N_STEPS = 300\n", "\n", "solver = optax.lbfgs(\n", " memory_size=20,\n", " linesearch=optax.scale_by_zoom_linesearch(max_linesearch_steps=10),\n", ")\n", "\n", "\n", "@jax.jit\n", "def step(params, opt_state, lam):\n", " def loss_fn(p):\n", " return loss(p, lam)\n", "\n", " value, grad = optax.value_and_grad_from_state(loss_fn)(params, state=opt_state)\n", " updates, opt_state = solver.update(grad, opt_state, params, value=value, grad=grad, value_fn=loss_fn)\n", " return optax.apply_updates(params, updates), opt_state\n", "\n", "\n", "def reconstruct(lam):\n", " params, opt_state = params_init, solver.init(params_init)\n", " for _ in range(N_STEPS):\n", " params, opt_state = step(params, opt_state, lam)\n", " return params" ] }, { "cell_type": "markdown", "id": "ff4afd7a", "metadata": {}, "source": [ "## How much regularization?\n", "\n", "We reconstruct at three weights, starting each run from the same random initial guess:\n", "$\\lambda = 0$ with the prior off, a small $\\lambda = 10^{-3}$, and a near-optimal\n", "$\\lambda = 0.3$. For each we record the relative field error\n", "$\\mathcal{L}_\\mathrm{data}$, the energy $\\mathcal{L}_\\mathrm{energy}$, the converged\n", "distance $d_\\mathrm{NV}^\\ast$, and, because this is a synthetic problem where the answer\n", "is known, the actual error of the reconstruction.\n", "\n", "On real data the true error does not exist, so choosing $\\lambda$ takes an independent\n", "argument: the corner of an L-curve traced by a full sweep over $\\lambda$, a noise\n", "estimate, or cross-validation on held-out pixels. The optimum also moves with the\n", "sample, so it is not a constant to carry between experiments." ] }, { "cell_type": "code", "execution_count": 12, "id": "b3cec459", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:36:10.788125Z", "iopub.status.busy": "2026-09-07T08:36:10.787448Z", "iopub.status.idle": "2026-09-07T08:38:23.472621Z", "shell.execute_reply": "2026-09-07T08:38:23.471325Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "ground truth: L_data = 0.0989 L_energy = -1.755\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "lambda = 0 L_data = 0.0945 L_energy = +8.276 d_NV = 24.88 nm RMS error in m = 1.202\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "lambda = 0.001 L_data = 0.0978 L_energy = -1.647 d_NV = 39.83 nm RMS error in m = 0.168\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "lambda = 0.3 L_data = 0.0984 L_energy = -1.755 d_NV = 40.00 nm RMS error in m = 0.013\n" ] } ], "source": [ "LAM_SMALL = 1e-3\n", "LAM_OPT = 0.3 # near-optimal weight, found by a sweep over lambda\n", "LAMBDAS = [0.0, LAM_SMALL, LAM_OPT]\n", "\n", "print(f\"ground truth: L_data = {float(L_data(m_true, D_NV_TRUE)):.4f} L_energy = {float(L_energy(m_true)):+.3f}\")\n", "\n", "results = {}\n", "for lam in LAMBDAS:\n", " params = reconstruct(lam)\n", " m = to_magnetization(params)\n", " d_nv = 1e-9 * float(jnp.abs(params[\"d_nv_nm\"]))\n", " h_fit = forward(m, d_nv)\n", " results[lam] = {\n", " \"m\": np.asarray(m[:, :, 0, :]),\n", " \"h\": np.asarray(h_fit),\n", " \"d_nv_nm\": 1e9 * d_nv,\n", " \"l_data\": float(jnp.linalg.norm(h_fit - h_meas) / H_MEAS_NORM),\n", " \"l_energy\": float(L_energy(m)),\n", " \"error\": float(jnp.sqrt(jnp.mean(jnp.sum((m[:, :, 0, :] - m_true[:, :, 0, :]) ** 2, -1)))),\n", " }\n", " print(\n", " f\"lambda = {lam:<6g} L_data = {results[lam]['l_data']:.4f}\"\n", " f\" L_energy = {results[lam]['l_energy']:+.3f}\"\n", " f\" d_NV = {results[lam]['d_nv_nm']:5.2f} nm\"\n", " f\" RMS error in m = {results[lam]['error']:.3f}\"\n", " )" ] }, { "cell_type": "markdown", "id": "1289ea59", "metadata": {}, "source": [ "## What the three runs show\n", "\n", "With the prior off ($\\lambda = 0$) the fit reproduces the measurement and still fails:\n", "at this distance the data leave most degrees of freedom unconstrained, and the optimizer\n", "spends them fitting noise. The converged distance drifts to 25 nm, because lowering the\n", "distance sharpens the operator and the noise-fitting texture prefers it.\n", "\n", "The regularized runs recover the true 40 nm from a start 15 nm below, to within 0.2 nm\n", "at both weights. At $\\lambda = 10^{-3}$ the labyrinth is already recovered at an RMS\n", "error of 0.17, overlaid with a short-wavelength ripple that the damping\n", "$e^{-k d_\\mathrm{NV}}$ hides from the scan. At $\\lambda = 0.3$ the ripple is gone and\n", "the RMS error drops to 0.01, at a data misfit indistinguishable from the weakly\n", "regularized run." ] }, { "cell_type": "markdown", "id": "9b77e127", "metadata": {}, "source": [ "The last column shows what each reconstruction looks like to the\n", "instrument: the field its magnetization produces at its converged $d_\\mathrm{NV}^\\ast$,\n", "with the measurement itself in the top row. The rows below it step through the\n", "regularization: none, small, near-optimal. The $\\lambda = 0$ row reproduces the\n", "measurement while its magnetization is speckled: the data alone does not pin down the\n", "texture completely, and the prior supplies the difference. At $\\lambda = 10^{-3}$ the\n", "labyrinth is already recovered, with the residual ripple visible in all three\n", "components. At $\\lambda = 0.3$ both the field and the magnetization match." ] }, { "cell_type": "code", "execution_count": 13, "id": "c733dc87", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:38:23.475430Z", "iopub.status.busy": "2026-09-07T08:38:23.474720Z", "iopub.status.idle": "2026-09-07T08:38:23.709205Z", "shell.execute_reply": "2026-09-07T08:38:23.708056Z" } }, "outputs": [], "source": [ "rows = [(\"ground truth\", np.asarray(m_true[:, :, 0, :]), np.asarray(h_meas), 1e9 * D_NV_TRUE)]\n", "rows += [(rf\"$\\lambda$ = {lam:g}\", results[lam][\"m\"], results[lam][\"h\"], results[lam][\"d_nv_nm\"]) for lam in LAMBDAS]\n", "\n", "h_vmax = 1e-3 * max(np.abs(h).max() for _, _, h, _ in rows)\n", "fig, axes = plt.subplots(len(rows), 4, figsize=(12.5, 3.1 * len(rows)), constrained_layout=True)\n", "for i, (axrow, (label, m, h, d_nv_nm)) in enumerate(zip(axes, rows)):\n", " im = plot_m(m, label, axrow[:3], titles=(i == 0))\n", " im_h = axrow[3].imshow(1e-3 * h.T, cmap=\"RdBu_r\", origin=\"lower\", vmin=-h_vmax, vmax=h_vmax)\n", " axrow[3].set_xticks([])\n", " axrow[3].set_yticks([])\n", " axrow[3].set_xlabel(rf\"$d_\\mathrm{{NV}} = {d_nv_nm:.1f}$ nm\")\n", " if i == 0:\n", " axrow[3].set_title(r\"$H^\\mathrm{dem}_\\mathrm{NV}$ (kA/m)\")\n", "fig.colorbar(im, ax=axes[:, :3], shrink=0.5, pad=0.01)\n", "fig.colorbar(im_h, ax=axes[:, 3:], shrink=0.5, pad=0.02)\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "92da0699", "metadata": {}, "source": [ "## Write the result\n", "\n", "`write_vti` stores the reconstruction for inspection in ParaView, as cell data alongside\n", "the material density that marks which cells are film and which are vacuum." ] }, { "cell_type": "code", "execution_count": 14, "id": "b425fbca", "metadata": { "execution": { "iopub.execute_input": "2026-09-07T08:38:23.711064Z", "iopub.status.busy": "2026-09-07T08:38:23.710872Z", "iopub.status.idle": "2026-09-07T08:38:23.984577Z", "shell.execute_reply": "2026-09-07T08:38:23.983327Z" } }, "outputs": [], "source": [ "state.m.tensor = jnp.zeros((N_XY, N_XY, NZ, 3)).at[:, :, 0, :].set(jnp.asarray(results[LAM_OPT][\"m\"]))\n", "state.write_vti([\"m\", \"rho\"], \"magnetization_reconstruction_from_nv_height_fit/m_reconstructed.vti\")" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.12" } }, "nbformat": 4, "nbformat_minor": 5 }