MagneX: A High-Performance, GPU-Enabled, Data-Driven Micromagnetics Solver for Spintronics

TL;DR

MagneX introduces a GPU-accelerated, AMReX-based micromagnetics solver that integrates multirate time integration (MRI) via SUNDIALS and a data-driven demagnetization surrogate using a Fourier Neural Operator. The framework supports detailed physics including Zeeman, demagnetization, anisotropy, exchange, and DMI, with a modular architecture that can swap the demagnetization calculation for a neural surrogate, all validated against μMAG standard problems and skyrmion benchmarks. Key findings show that MRI can reduce time-to-solution by about 48% compared to RK4 in stiff regimes while preserving accuracy, and that NN-based demagnetization can reproduce essential dynamical features with high fidelity. The work demonstrates scalable performance on modern HPC systems and offers open-source tools for integrating high-fidelity micromagnetic simulations with data-driven surrogates, enabling efficient exploration of multiphysics spintronic devices.

Abstract

In order to comprehensively investigate the multiphysics coupling in spintronic devices, it is essential to parallelize and utilize GPU-acceleration to address the spatial and temporal disparities inherent in the relevant physics. Additionally, the use of cutting-edge time integration libraries as well as machine learning (ML) approaches to replace and potentially accelerate expensive computational routines are attractive capabilities to enhance modeling capabilities moving forward. Leveraging the Exascale Computing Project software framework AMReX, as well as SUNDIALS time-integration libraries and python-based ML workflows, we have developed an open-source micromagnetics modeling tool called MagneX. This tool incorporates various crucial magnetic coupling mechanisms, including Zeeman coupling, demagnetization coupling, crystalline anisotropy interaction, exchange coupling, and Dzyaloshinskii-Moriya interaction (DMI) coupling. We demonstrate the GPU performance and scalability of the code and rigorously validate MagneX's functionality using the mumag standard problems and widely-accepted DMI benchmarks. Furthermore, we demonstrate the data-driven capability of MagneX by replacing the computationally-expensive demagnetization physics with neural network libraries trained from our simulation data. With the capacity to explore complete physical interactions, this innovative approach offers a promising pathway to better understand and develop fully integrated spintronic and electronic systems.

Figures (11)

• Figure 1: Architecture of the MagneX micromagnetic solver. The Landau–Lifshitz–Gilbert (LLG) equation is evolved using a multirate time integration scheme (SUNDIALS), which treats fast (exchange, anisotropy, DMI) and slow (demagnetization) components on separate time scales. Spatial discretization and parallelization are handled via the AMReX framework. The demagnetizing field ${\bf{H}_{\rm demag}}$ can be computed either via FFT-based spectral convolution or approximated by a neural network surrogate. This hybrid design enables efficient, scalable, and extensible simulations on modern multi-GPU architectures.
• Figure 2: Machine learning pipeline for demagnetization in MagneX. A supervised learning framework based on PyTorch is used to train a Fourier Neural Operator (FNO) model to approximate the demagnetizing field, ${\bf{H}_{\rm demag}}$, from the magnetization field, $\mathbf{M}$. The trained model is JIT-traced for C++ compatibility and integrated into the MagneX simulation workflow. The pipeline includes data normalization, tensor conversion between column-major (Fortran) and row-major (C++/Python) memory layouts, and bidirectional index mapping between AMReX’s Multifab data structures and PyTorch tensors. This enables efficient, runtime evaluation of learned surrogates during micromagnetic simulations.
• Figure 3: Validation of MagneX against $\mu$MAG Standard Problem 2. (Top) Geometry of the thin ferromagnetic prism used in the benchmark, with aspect ratios $L/d = 5.0$ and $t/d = 0.1$, as specified by the problem. The next three panels show a comparison of MagneX results with reference data for the normalized coercive field, and $m_x$ and $m_y$ remanent magnetization components as functions of normalized size $d/\ell_\mathrm{ex}$. The benchmark examines the transition from coherent rotation to vortex-mediated reversal as $d/\ell_\mathrm{ex}$ increases. MagneX results show strong agreement with previously submitted solutions hosted on the $\mu$MAG website.
• Figure 4: ${\bf{M}}$ vectors illustrating (top) flower and (bottom) vortex states generated at the critical domain size $L/l_{\rm ex} = 40$.
• Figure 5: Images of magnetization when $m_x$ first crosses zero and evolution of $(m_x, m_y, m_z)$ for (top) Field 1 and (bottom) Field 2 in Standard Problem 4 at both coarse and fine resolution, and comparison with McMichael et al. mcmichael2001switching$\mu$MAG website results.