Content area
We present NeuralMag, a flexible and high-performance open-source Python library for micromagnetic simulations. NeuralMag leverages modern machine learning frameworks, such as PyTorch and JAX, to perform efficient tensor operations on various parallel hardware, including CPUs, GPUs, and TPUs. The library implements a novel nodal finite-difference discretization scheme that provides improved accuracy over traditional finite-difference methods without increasing computational complexity. NeuralMag is particularly well-suited for solving inverse problems, especially those with time-dependent objectives, thanks to its automatic differentiation capabilities. Performance benchmarks show that NeuralMag is competitive with state-of-the-art simulation codes while offering enhanced flexibility through its Python interface and integration with high-level computational backends.
Introduction
Micromagnetic simulations are a fundamental tool in the study of magnetization dynamics and play a crucial role in understanding and designing magnetic materials and devices. These simulations model the behavior of magnetic and magnonic systems at the nanoscale, providing insight into phenomena such as domain wall motion, magnetization reversal, and spin wave propagation. The field relies on various computational methods, with finite-difference and finite-element schemes being widely used. Notable examples of established finite-difference codes include OOMMF1 and fidimag2 for CPU-based simulations and mumax33, BORIS4, and magnum.np5 for GPU-accelerated simulations. Finite-element-based methods, such as those implemented in NMag6, Tetramag7, FastMag8, FinMag9, and magnum.fe10, provide greater flexibility in handling complex geometries but can be computationally more expensive. More recently, the finite-element solver TetraX11 has gained popularity in the magnonics community due to its efficient eigenmode solver in infinite geometries.
In addition to standard micromagnetic simulations, inverse problems have attracted considerable attention in recent years. These problems involve determining the optimal parameters, such as material properties, external fields, or device geometries, that lead to a desired magnetic configuration or device functionality. A significant body of work has focused on inverse modeling of the demagnetization field, a static inverse problem. This has been particularly useful in the context of magnetic 3D printing, where topology optimization techniques are employed to design optimal material layouts, and the inverse modeling is used to infer the magnetization configuration of printed samples12, 13–14.
More recently, research in the emerging field of inverse magnonics has gained momentum, focusing on optimizing the functionality of magnonic devices. Magnonics uses spin waves (magnons) for information processing, and designing efficient magnonic devices poses complex nonlinear optimization challenges. Inverse-design approaches have been increasingly applied to magnonics, allowing researchers to automate the design of devices by specifying a desired functionality and using computational algorithms to find the optimal configuration15, 16, 17–18.
In this paper, we present a novel discretization strategy for micromagnetic simulations, adjoint-state algorithms for efficiently solving time-dependent inverse problems, and the software design of NeuralMag, which integrates these advancements into a flexible and high-performance computational framework.
Results
Micromagnetics
The micromagnetic model provides a semi-classical continuum description of magnetization dynamics in ferromagnetic systems, as originally formulated by Brown19. The key governing equation is the Landau–Lifshitz–Gilbert (LLG) equation, which reads
1
with m being the unit-vector field representation of the magnetization, γ being the reduced gyromagnetic ratio, and α being a dimensionless damping parameter. The effective field Heff accounts for all relevant interactions within the system and derives from the total energy as2
with Ms being the saturation magnetization and δE/δm denoting the variational derivative of the energy with respect to the magnetization20,21. When the energy E depends on spatial derivatives of the magnetization field m, additional boundary conditions must be imposed to solve Eqs. (1) and (2). One such example is the micromagnetic exchange energy, which is defined as3
where A is the exchange stiffness constant. The variation of the exchange energy with respect to m yields4
leading to the exchange field definition5
The boundary term in Eq. (4) defines the appropriate exchange boundary condition. To satisfy equilibrium conditions in micromagnetics, the system must fulfill Brown’s conditions, which require m × δE/δm = 0 for x ∈ Ω, and m × B = 0 for x ∈ ∂Ω.
A similar variational treatment at internal interfaces, where material parameters vary discontinuously, introduces additional interface conditions22. Assuming a continuous magnetization across such interfaces and dividing the domain into regions of continuous material parameters, the corresponding interface condition can be written as m1 × B1(n) = m2 × B2(n), where B1(n) and B2(n) represent the boundary terms on either side of the interface.
In case of the exchange energy being the only energy contribution introducing spatial derivatives and furthermore considering m⊥∂m/∂n, this leads to the well-known exchange jump condition19
6
In addition to satisfying equilibrium conditions, the boundary and interface conditions must be consistently fulfilled at all times when solving the LLG equation22.NeuralMag implements a novel nodal finite-difference scheme described in section “Nodal finite-difference scheme” to accurately solve the micromagnetic equations. The code supports both PyTorch23 and JAX24 as computational backends, enabling efficient tensor operations and automatic differentiation on various hardware platforms. Details on the implementation are provided in section “Implementation”.
Inverse micromagnetics
In addition to employing a nodal finite-difference scheme, NeuralMag is specifically designed to address inverse problems in both space and time domains. In this context, the computation of individual field terms or the solution of the LLG Eq. (1) is classified as a forward problem F. Given a vector of design variables θ, which may include material properties or the initial magnetization configuration, these forward problems yield well-defined outputs y, such as effective field contributions or the resulting magnetization trajectory
7
An inverse problem is formulated to determine the design variables θ that yield a specified result y from the forward problem. This task is often challenging, as inverse problems are typically ill-posed, and their solution vectors may encompass a large number of degrees of freedom. The most common strategy to solve such a problem is the reformulation in terms of a minimization problem that might be complemented by additional terms for regularization or smoothing purposes. In the case of a high-dimensional input θ and a nonlinear function F, this problem is nontrivial. In such cases, iterative methods, typically based on the gradient of the functional , are commonly employed to find a solution. NeuralMag uses automatic differentiation25 for static problems such as inverse strayfield calculations. In contrast to the adjoint method that has been used in previous works12,13, this approach performs the differentiation on the discrete level (discretize first). As for the adjoint method, the gradient computation requires a forward solve and a subsequent backward solve with the complexity of the backward solve being equivalent to that of the forward solve.8
For time-dependent problems, NeuralMag implements the adjoint-state method26. The adjoint-state method is a powerful tool for the solution of PDE-constrained optimization problems, also referred to as optimal-control problems. Given a forward problem
9
with design variables θ, we define an objective functional10
with m(T; θ) being the solution of Eq. (9) for a final time T and ytarget being the desired output of the forward problem. In order to compute the gradient of the objective functional with respect to the design variables , the adjoint-state method requires two steps. In the first step, the forward problem (9) is solved for the given design variables θ, which results in the output moutput = m(T). In the second step, the so-called adjoint problem is solved, which is given by the following system of ODEs11
with a being the so-called adjoint variable. This system is solved backwards in time, starting from the final time T used in the forward pass. Successful integration of the system yields the output u(0), which can be identified as the desired gradient of the objective12
While the objective (10) depends solely on the magnetization at the final time T, extending this method to objectives depending on multiple time points Ti can be done in a straightforward fashion by adding appropriate terms depending on m(Ti; θ) to (10). The computational and storage complexity of the adjoint system is comparable to that of the forward problem, yielding an exceptionally efficient strategy for the gradient computation of PDE-constrained optimization. This method is superior to the backpropagation method15,16 with regard to the storage requirements that are similar to a regular forward pass. However, this advantage comes at the cost of reduced accuracy, which is caused by the backwards pass that reconstructs the magnetization trajectory by inverse integration instead of using the exact values from the forward pass.Validation and benchmarks
To validate the accuracy of NeuralMag, we solve two significant micromagnetic problems. These tests showcase NeuralMag’s ability to handle both standard and advanced cases, verifying its precision and computational efficiency.
The first validation case is MuMag Standard Problem #427, which simulates the dynamic behavior of a thin ferromagnetic film under an applied magnetic field tilted either by 170∘ or 190∘ toward the x-axis. The focus is on the time evolution of the averaged magnetization components. We solve this problem using a full 3D spatial discretization and compare the results to a 2D simulation as described in section “Low-Dimensional Geometries” of the paper. The results for the field tilted by 170∘, displayed in Fig. 1, show excellent agreement with the reference solutions from the MuMag community, demonstrating the precision of NeuralMag in simulating the time dynamics of micromagnetic systems both with the 3D as well as 2D thin-film approximation. The second part of the standard problem #4, which simulates the switching under a field tilted by 190∘ is included in the demos that are accessible via the NeuralMag website28.
[See PDF for image]
Fig. 1
MuMag standard problem #4.
The results are presented using both 2D and 3D discretizations as computed by NeuralMag. The reference solution, computed with OOMMF1, is depicted by solid lines for comparison. The NeuralMag solutions are illustrated using circles for the 3D discretization and squares for the 2D discretization.
The second validation case involves solving the domain wall pinning problem proposed by Heistracher et al.29. This problem focuses on calculating the coercive field required to unpin a domain wall at the interface between two magnetic phases with varying material properties, such as exchange interaction, uniaxial anisotropy, and spontaneous magnetization. This problem is sensitive to discontinuities in these parameters, making it an ideal test for NeuralMag’s handling of complex material boundaries.
In this validation, we compare the switching fields calculated by NeuralMag with the analytical results provided in Table 1 of the original paper. We varied the material parameters (exchange constant A, anisotropy constant K, and saturation magnetization Ms) in different combinations across the two magnetic phases. Table 1 compares the switching fields obtained using NeuralMag with those presented in the reference paper. Our results closely match the analytical solutions, with minor deviations likely due to the time integration method and field rate used during the simulation. These successful validations confirm that NeuralMag correctly handles discontinuities at material interfaces and provides accurate predictions for complex micromagnetic systems.
Table 1. Depinning fields for a domain wall in a two-phase magnet, as defined in ref.29, computed with NeuralMag and compared to the analytical and numerical reference solutions computed with magnum.af
Discontinuous parameters | Analytical [T] | magnum.af [T] | NeuralMag [T] |
|---|---|---|---|
A/K/Ms | 1.568 | 1.585 | 1.580 |
A/K | 1.089 | 1.116 | 1.112 |
A/Ms | 1.206 | 1.256 | 1.205 |
A | 0.838 | 0.868 | 0.867 |
K/Ms | 1.005 | 1.020 | 1.012 |
K | 0.565 | 0.582 | 0.571 |
To evaluate the performance of NeuralMag, we conducted a throughput benchmark, shown in Fig. 2, where we compare the time required for evaluating the right-hand side (RHS) of the LLG equation across different system sizes. Specifically, we measure the time for the integration of the full LLG, including the exchange and demagnetization field, and then divide by the number of field evaluations. This procedure can be easily applied to any micromagnetic code without the need to modify it and provides a robust measure of the overall performance at the same time. In this benchmark, NeuralMag is compared to two widely-used micromagnetic simulation tools: mumax33 and magnum.np5. mumax3 shows the best performance due to its highly optimized GPU implementation. However, NeuralMag, when using JAX as the backend, almost matches the performance of mumax3, being less than a factor of 2 slower.
[See PDF for image]
Fig. 2
Computational benchmark.
Comparison of the computation time for evaluating the right-hand side of the Landau-Lifshitz-Gilbert (LLG) equation, including both the exchange and demagnetization fields, across various system sizes N, with NeuralMag (NM) in comparison with other finite-difference codes. The legend indicates the code as well as the device that was used for the computation. As a GPU, we used an NVIDIA A100 card with 80 GB of RAM, and as a CPU, we used an Intel Xeon Gold 6326. For GPU timings, solid lines indicate single-precision computations while dashed lines indicate double-precision. CPU timings were all performed with double precision, with solid lines indicating single-core computations and dashed lines indicating the use of 8 cores.
Remarkably, NeuralMag maintains this competitive performance even for small system sizes, despite the computational overhead typically associated with a Python implementation. This performance can be attributed to the just-in-time (JIT) compilation feature of JAX, which optimizes the entire RHS of the LLG equation at runtime. Thanks to NeuralMag’s architecture, JAX is able to analyze and compile the full computation into highly optimized machine code, reducing overhead and achieving near-optimal execution times. This demonstrates the strength of NeuralMag’s design in leveraging modern machine learning frameworks to achieve high-performance computations while maintaining flexibility.
The remaining performance gap of approximately a factor of two compared to MuMax3 likely arises from the demagnetization field computation, as the current FFT interface in JAX is limited, preventing certain optimizations. However, as JAX’s FFT capabilities expand, this gap could narrow significantly—or even vanish completely—in the future.
The CPU performance using the JAX backend is approximately 10 times slower than the single-precision GPU performance, which is remarkably fast. However, there is no notable difference in timings for single CPU usage compared to the OpenMP parallelization on 8 CPUs. For very small systems, the CPU implementation of JAX even outperforms the GPU implementation. This is most likely due to the missing overhead of kernel deployment and data transfer to the GPU, which leads to almost linear scaling down to very small system sizes. CPU computations with PyTorch are approximately 10 times slower than JAX computations.
Because the adjoint/autodiff gradient has the same asymptotic cost as one forward RHS evaluation, the timings in Fig. 2 also characterise a single gradient step in inverse-design workflows.
Inverse problems
Listing 1. Gradient computation of a topology optimization problem.
A first, purely static inverse problem concerns the topology optimisation of a hard-magnetic cuboid of dimensions 100 nm × 100 nm × 50 nm that is uniformly magnetised along +z. Using the Solid-Isotropic Material with Penalisation (SIMP) approach30, the material density field ρ(x) ∈ [0, 1] is restricted to the magnetic design region Ωm, see Listing 1). The objective
13
with the density14
maximises the z-component of the stray field at the probe point xcenter located 5 nm above the centre of the top surface. Because the demagnetization field in NeuralMag is differentiable with respect to ρ, its gradient is obtained in a single reverse pass and fed to any gradient-based optimiser that enforces 0 ≤ ρ ≤ 1. The relevant part of the simulation script is shown in Listing 1. The final topology is shown in Fig. 3 in excellent agreement with published solutions5,14.[See PDF for image]
Fig. 3
Topology optimization.
Optimization result of a permanent magnetic sample with magnetization M = (0, 0, Mz) that maximizes Hz(M) at a single point above the sample marked by the grey sphere.
Listing 2. Simulation script for inverse problem.
After the static topology-optimisation example above, we now turn to a genuinely dynamic inverse problem. Specifically, we aim to optimize the direction of an external magnetic field to align the magnetization of a single-domain particle with a target configuration mtarget after a given time T, see Fig. 4(a). The optimization minimizes the objective function with respect to the field angles θ and ϕ, as defined by the system
15
16
17
with m(t) being constrained by the LLG (1). A shortened code listing demonstrating the setup for this inverse problem is provided in Lst. 2. In this simple optimization, convergence is achieved after 30–50 gradient-descent steps, see Fig. 4b. NeuralMag computes the gradient of the objective function by performing one forward and one backward simulation of the dynamic problem.[See PDF for image]
Fig. 4
Dynamic inverse problem.
Simple inverse micromagnetic problem for the optimization of the external field direction in order to align the magnetization of a single-domain particle in a given direction. a Sketch of the problem setup. b Convergence of the objective function and the optimized field angles θ and ϕ.
Finally, a recently published application study by Voronov et al.31 demonstrates that NeuralMag also scales to highly complex, fully dynamic topology-magnetization tasks: the framework tackled (i) the inverse design of a Stoner–Wohlfarth nanoparticle that converged in ≈250 adjoint iterations and (ii) a 1 μm × 1 μm magnonic demultiplexer, parameterised by 400 radial-basis functions on a 512 × 64 × 1 mesh, which achieved more than an order-of-magnitude spin-wave-contrast between 2.6 GHz and 2.8 GHz channels after only ≈100 optimization steps—underscoring the robustness and scalability of our approach for real-world, nonlinear device geometries.
Methods
Nodal finite-difference scheme
Existing micromagnetic simulation software usually employs either a finite-difference discretization on regular grids22,32 or a finite-element discretization on irregular grids20,22. The use of regular cuboid grids in the case of finite-difference micromagnetics allows for a very efficient computation of the demagnetization field by means of an FFT-accelerated convolution. On the other hand, the finite-element method allows for the accurate modeling of complex structures due to the use of irregular meshes.
Moreover, finite-element micromagnetics provides a more subtle but sometimes highly relevant advantage over finite-difference micromagnetics: In finite-element micromagnetics, the magnetization is usually explicitly defined on each mesh-vertex, whereas standard finite-difference tools store one magnetization vector per simulation cell, which is typically taken to be the magnetization in the center of this cell. While this difference appears to be insignificant for the micromagnetic modeling in the bulk, it plays a crucial role when considering material interfaces where the magnetization is subject to boundary and jump conditions. Consider e.g., the exchange jump condition (6), which prescribes a discontinuity in the first spatial derivative of the magnetization across material interfaces. Choosing the degrees of freedom of the magnetization in the cell centers, as illustrated in Fig. 5a, requires a careful treatment of the boundary conditions that are defined on the vertices29. Similar considerations apply to interfacial energy contributions such as the RKKY coupling between two ferromagnetic layers33. Inaccurate modeling of the boundary conditions can lead to a loss of convergence order and consequently, introduce significant numerical errors. Introducing more energy contributions depending on surface integrals or spatial derivatives of m results in more complex boundary conditions22 that become unfeasible to handle in standard finite-difference micromagnetics.
[See PDF for image]
Fig. 5
Discretization strategies.
Illustration of the discretization of the magnetization m and the material parameter A for a one-dimensional representation of a two-phase magnetic system, using different numerical schemes: a Standard finite differences: Both the material parameter and the magnetization are discretized with a single value per simulation cell. The magnetization degrees of freedom are treated as sample points of a continuous function. b Finite elements: The material parameters are discretized using piecewise constant functions, while the magnetization is represented as piecewise affine, with degrees of freedom located at the vertices.
In contrast, the finite-element method allows for the choice of tailored function spaces for the magnetization and material parameters, as shown in Fig. 5b. Moreover, the inherently variational nature of the finite-element method allows to solve for the effective-field contributions by directly considering the variation of the energy22, resulting in the weak form
18
with V being a sufficiently smooth function space referred to as the test space. By a proper choice of function spaces for the material parameters and fields, this procedure does not require explicitly accounting for the boundary conditions at all.Local field terms
The nodal finite-difference scheme proposed in this work applies the finite-element method for local field contributions on a regular cuboid grid. This enables the use of an FFT-accelerated demagnetization-field computation as in standard finite-difference micromagnetics, see section “Demagnetization Field”, while providing the rigorous and accurate handling of material interfaces for all local field contributions due to finite-element modeling. In order to address the cells and nodes of the regular grid, we introduce multi-indices c, n, and i as
19
20
21
with N1, N2, and N3 being the number of simulation cells in the respective mesh dimension. The indices c and n are used to address simulation cells and nodes, respectively, according to the numbering introduced in Fig. 6a. The index i either acts as a local vertex number in a simulation cell according to Fig. 6b or, more generally, as a relative index to address neighborships.[See PDF for image]
Fig. 6
Numbering of degrees of freedom.
Cell and vertex numbering using multiindices for nodal finite differences. a Two-dimensional representation of global cell and node indices denoted by c (black) and n (blue). b Three-dimensional representation of local vertex numbering denoted by index i (blue).
We discretize all continuous fields appearing in weak forms with standard piecewise polynomial and globally continuous basis functions ϕn that form a nodal basis on the cuboid mesh. Each basis function ϕn is defined per simulation cell in terms of reference basis functions as
22
with Δxk being the simulation-cell size in the k-th dimension. The reference basis functions are defined on the reference unit cell Ωref = [0, 1] × [0, 1] × [0, 1] as23
where denotes the characteristic function of Ωref, which evaluates to 1 if x ∈ Ωref and to 0 else. This restricts the support of the reference basis functions to the reference cell Ωref. A 2D representation of a basis function is visualized in Fig. 7. Furthermore, we introduce vector basis functions as24
with ej being the unit vector in direction j ∈ {1, 2, 3}. Continuous vector fields such as the magnetization m and the effective field Heff are then discretized as25
with the superscript h denoting the discretized version of a field and coefficients mn,j being the nodal values of the vector field m.[See PDF for image]
Fig. 7
Nodal basis function.
Two-dimensional representation of a basis function ϕn in nodal finite differences with a support spanning 4 simulation cells.
For material parameters, such as the saturation magnetization Ms, we choose a piecewise constant function space in order to allow for the accurate modeling of rapid material interfaces. Namely, we define these parameters per simulation cell resulting in the following discretization
26
with basis functions27
Replacing all fields with their discretized counterparts in the weak form (18) and testing with individual basis functions instead of arbitrary test functions yields the discretized weak form28
For a given node n, we split the variation δE(mh, ϕn,j) into its contributions from the eight simulation cells that share node n and we address these cells by the local index i ∈ {0, 1}3. In general, the variation over a single simulation cell depends on the magnetization values of all nodes of this cell. Considering the three components of the magnetization, the contribution of the cell i to the variation can be written as
29
where denotes all nodal values of the magnetization in cell i. If the energy E is quadratic in m, the function F is linear in and can be described by a 24 × 24 matrix considering the 23⋅3 degrees of freedom defined by the index pairs i, j and . In the finite-element context, this matrix is usually referred to as the element matrix of the weak form.In order to compute the variation at all nodes, we introduce the vector δE with components δEn,j = δE(mh, ϕn,j) and the auxiliary vectors δE*i containing the cell-wise variations according to (29) for all nodes. Considering the node and cell numbering introduced in Fig. 6, the global node index is given by the global cell index and the relative node index as n(c, i) = c + i resulting in
30
31
Note that F only depends on the relative index i and the component j. Eqs. (30) and (31) deliver a straight-forward strategy for a parallel evaluation over the cell index c, see section “Form compilation”.If the energy E depends on further fields, such as an external field or material parameters, the mapping function can be easily extended by adding additional arguments
32
where the variables are the coefficients of arbitrary scalar fields discretized with nodal basis functions (22) and the variables are the coefficients of arbitrary scalar fields discretized with cell basis functions (27). Since Fj,k does not explicitly depend on the cell index i, it is fully determined by the integrand of the weak form (28) and the dimensions of a single simulation cell Ωi.In order to determine the discretized effective field Hh, the weak form requires the solution of a linear mass system defined by the left-hand side of Eq. (28). To avoid this costly procedure, we employ mass lumping to the left-hand side of Eq. (18) as described in Abert22 resulting in
33
where the saturation magnetization Ms is discretized cell-wise according to Eq. (26).The proposed method is applicable to any energy contribution whose density depends solely on the magnetization and its first-order spatial derivatives, such as Zeeman energy, crystalline anisotropies, and both symmetric and antisymmetric exchange interactions. Due to the regularity of the cuboidal grid, a matrix-free implementation of the presented scheme is straightforward. The local support of the basis functions results in a computational complexity of for the evaluation of any local field term, with N being the number of simulation cells.
Demagnetization field
To compute the demagnetization field, we employ the well-established FFT-accelerated method commonly used in standard finite-difference micromagnetic simulations32. This algorithm calculates the demagnetization field generated by homogeneously magnetized cuboidal simulation cells arranged on a regular grid through fast convolution. Since this method requires both the magnetization and the resulting field to be specified for each simulation cell, we introduce a straightforward pre- and post-processing step. This procedure averages the values to transition between nodal and cell-centered discretizations efficiently. FFT-accelerated methods that operate directly on node-wise discretized magnetizations have been proposed in previous studies34,35. However, we opt for the standard method based on homogeneously magnetized cuboids due to its advantages in memory efficiency and computational performance, specifically for 2D computations where the FFT also reduces to two dimensions.
Low-dimensional geometries
Discretizing a mesh with N1 × N2 × N3 cells results in (N1 + 1) × (N2 + 1) × (N3 + 1) degrees of freedom when using a nodal basis for the function discretization. In bulk system simulations, this introduces only a negligible overhead in comparison to standard finite-difference schemes, where the degrees of freedom are equal to the number of simulation cells. However, a significant application area for micromagnetic simulations involves magnetic thin films, which are often discretized with just a single layer of simulation cells. In such cases, the 3D nodal discretization introduces a notable overhead, roughly doubling the computational cost, because it requires separate descriptions for the top and bottom surfaces of the thin film. This contrasts with standard finite differences, where the problem effectively reduces to a 2D formulation. By transitioning to 2D basis functions while maintaining full 3D integration in the weak form (28), the nodal finite-difference scheme can accurately describe magnetic thin films. This approach reduces the degrees of freedom to (N1 + 1) × (N2 + 1) × 1, making it more efficient for thin film simulations. Namely, the 2D basis function on the reference cell are chosen as
34
with a 2D multiindex i = (i1, i2) ∈ {0, 1}2. As illustrated in Fig. 8 this approach can also be generalized to 1D problems, leading to basis functions35
with a scalar index i ∈ {0, 1}.[See PDF for image]
Fig. 8
Low-dimensional discretization.
Representation of the degrees of freedom for a square-shaped rod using the following methods: a full three-dimensional discretization, b two-dimensional discretization with basis functions that are constant along the third dimension, and c one-dimensional discretization with basis functions that are constant along both the second and third dimensions.
Implementation
NeuralMag is a Python library designed specifically for micromagnetic simulations, with a focus on high-performance tensor computations. A key feature of NeuralMag is its ability to operate with either PyTorch23 or JAX24,36 as a computational backend, allowing users to select the framework that best suits their needs. By leveraging these modern machine learning frameworks, NeuralMag achieves efficient computations on a variety of parallel hardware, including CPUs, GPUs, and TPUs. This versatility is complemented by the advantages these frameworks offer, such as optimized performance for large-scale simulations and built-in support for automatic differentiation, which simplifies solving inverse problems. Through the modular design, the software is prepared to simplify the use of other computational backends in the future.
The use of either PyTorch or JAX as backends allows NeuralMag to fully exploit the unique strengths of each framework. PyTorch’s
In contrast, JAX’s
Both backends support single- and double-precision computations, enabling NeuralMag to offer users flexibility in balancing computational speed with numerical accuracy according to the requirements of each simulation. The dual-backend approach ensures that NeuralMag can adapt to the user’s preferred ecosystem while maintaining high computational efficiency and flexibility.
Form compilation
Listing 3. Symbolic definition of the exchange energy (3) in NeuralMag.
Listing 4. Automatically generated code for the computation of the exchange field.
At the heart of NeuralMag is a form compiler that translates a symbolic representation of a finite-element weak form into efficient tensor operations tailored to the chosen backend. For symbolic computation, NeuralMag leverages the Python library SymPy37. SymPy provides a powerful framework for representing the mathematical structures involved in micromagnetic simulations. Specifically, NeuralMag introduces custom SymPy symbols to represent functions that are discretized either node-wise or cell-wise, as described in the section “Nodal finite-difference scheme” of this paper. Users can define the weak form of the micromagnetic problem using SymPy’s symbolic language, allowing them to work in an intuitive mathematical formulation.
In addition to defining weak forms symbolically, NeuralMag leverages SymPy to automatically perform the variation of a symbolic energy expression, allowing it to derive the corresponding weak form. This capability streamlines the process of converting complex energy functionals into their weak form representations. For instance, in Lst. 3, the exchange energy is defined symbolically using SymPy, demonstrating how users can express physical energy terms within the framework.
NeuralMag’s form compiler processes the symbolic weak form and transforms it into the discrete mapping function Fi,j, as defined in Eq. (30). This transformation is achieved by applying Gauss quadrature to integrate over the finite elements, converting the weak form into a set of tensor operations—primarily multiplications and summations—that can be efficiently executed by the selected backend. The role of the relative cell index, as discussed in Eqs. (30) and (31) are handled by tensor slicing. This involves slicing along specific tensor dimensions by removing either the first
Dynamic attributes
Listing 5. Example usage of dynamic attributes in NeuralMag.
Listing 6. Automatically generated function for the evaluation of the dynamic attribute
NeuralMag introduces the concept of dynamic attributes through its state object, which allows attributes to be either tensors or functions that depend on tensors and return tensors. This flexible design enables dynamic relationships between attributes, where some can be defined as functions of others, with NeuralMag automatically managing these dependencies. For example, consider the code in Lst. 5: attributes
When defining such dynamic attributes, NeuralMag analyzes the function signatures to identify all dependencies in a recursive manner. It then generates a new Python function at runtime that only relies on pure tensors and eliminates any control structures, such as loops or conditionals, ensuring the function remains optimal for high-performance tensor computation. In the case of the example from Lst. 5, the dynamically created function looks like Lst. 6, where
As an example, a material parameter such as the exchange constant A can be defined either as a regular, constant attribute or as a dynamic attribute, depending on other state variables such as the temperature and the time in order to simulate the magnetic response to a heat pulse. Either way, the value of A can be accessed by
Automatic differentiation and time integration
In the context of inverse problems, NeuralMag leverages automatic differentiation and efficient time integration to solve complex optimization tasks. Both PyTorch and JAX offer powerful automatic differentiation capabilities, which are crucial for computing gradients with respect to parameters in inverse problems. For time integration, NeuralMag integrates with torchdiffeq38 (for PyTorch) and diffrax39 (for JAX), both of which provide support for solving ordinary differential equations (ODEs).
Time integration is essential in dynamic micromagnetic problems, where the system’s evolution must be accurately tracked. Both libraries support a variety of numerical schemes for time stepping, including Euler methods, Runge–Kutta methods (such as RK4), and adaptive solvers like the Dormand–Prince method. These methods ensure that NeuralMag can flexibly adapt to different accuracy and performance requirements in dynamic simulations.
For gradient-based optimization in inverse problems, NeuralMag supports both the adjoint method26 and traditional backpropagation. The adjoint method is particularly well-suited for problems with long time horizons or large state spaces, as it computes gradients more efficiently by solving an adjoint ODE backward in time. Both torchdiffeq and diffrax support the adjoint method for time integration, offering an efficient way to compute gradients when optimizing over dynamic systems. At the same time, they also allow for direct backpropagation through the time integration process, which can be more straightforward for shorter time intervals or simpler problems.
By combining automatic differentiation with advanced time integration techniques, NeuralMag can effectively tackle inverse problems in micromagnetic simulations, allowing users to optimize parameters while ensuring accurate numerical solutions over time.
Discussion
In this paper, we have introduced NeuralMag, an open-source Python library for micromagnetic simulations that leverages modern machine learning frameworks such as PyTorch and JAX to achieve high performance. NeuralMag implements a novel nodal finite-difference discretization scheme, which provides a rigorous numerical description of continuous fields such as the magnetization as well as discontinuous material parameters. This approach is particularly useful for the accurate modeling of material interfaces while maintaining the same computational complexity as standard finite-difference schemes. Its performance is competitive with state-of-the-art micromagnetic simulation codes, yet it offers unparalleled flexibility due to its Python-based interface and support for optimized tensor operations on a variety of hardware platforms.
NeuralMag is especially well-suited for solving inverse problems, particularly those with time-dependent objectives, thanks to its ability to seamlessly compute gradients using automatic differentiation. This makes it a powerful tool for a wide range of optimization and simulation tasks in micromagnetics. NeuralMag is freely available28, making it accessible to the broader research community for further development and application.
Acknowledgements
This research was funded in whole or in part by the Austrian Science Fund (FWF) 10.55776/P34671, 10.55776/I6068, 10.55776/PAT3864023, and 10.55776/PIN1434524. For open access purposes, the author has applied a CC BY public copyright license to any author-accepted manuscript version arising from this submission. This work has been supported by the Horizon Europe research and innovation program through MaMMoS grant agreement No 101135546 and Marie Skłodowska Curie grant agreement No 101152613. We gratefully acknowledge the wedding of the Koraltans for fruitful discussions and a great time, which led to this publication.
Author contributions
C.A. conceived the project and developed the software architecture as well as the form compiler of NeuralMag. F.B. and R.A. contributed to coding efforts and performance optimization. R.K., P.F., T.S., and D.S. provided key input in the development of the numerical formalism. A.V., A.C., and S.K. contributed to the development and testing of the framework’s inverse problem-solving capabilities. S.A.P., S.H., M.L., and H.F. enhanced the codebase and established automatic code formatting and testing workflows. C.A. wrote the manuscript with contributions from all authors. All authors discussed the results and gave feedback on the manuscript.
Data availability
The code of NeuralMag is publicly available28. The datasets generated and analysed during the current study can be reproduced by running the respective demo scripts provided in the NeuralMag repository.
Code availability
The source code of NeuralMag is publicly available under the MIT License on GitLab at https://gitlab.com/neuralmag/neuralmag28. Comprehensive documentation, an API reference, and tutorials are available at https://neuralmag.gitlab.io/. NeuralMag can be installed via standard package managers such as pip or conda. Users and community members are encouraged to contribute to the codebase, tutorials, and documentation. Continuous integration workflows are set up using GitLab CI/CD to automatically run tests after every code change. These tests are run with both the PyTorch and the JAX backend and include unit, integration, and system tests, covering both the standard problems and benchmarks. The repository includes all numerical problems discussed in this paper, as well as the code to reproduce the benchmarks.
Competing interests
The authors declare no competing interests.
Publisher’s note Springer Nature remains neutral with regard to jurisdictional claims in published maps and institutional affiliations.
References
1. Donahue, M. J. & Porter, D. G. OOMMF User’s Guide, Version 1.0. Report No. 6376 https://doi.org/10.6028/NIST.IR.6376 (Accessed June 5, 2025) (National Institute of Standards and Technology, 1999).
2. Bisotti, M-A et al. Fidimag–a finite difference atomistic and micromagnetic simulation package. J. Open Res. Softw.; 2018; 6, 22. [DOI: https://dx.doi.org/10.5334/jors.223]
3. Vansteenkiste, A. et al. The design and verification of MuMax3. AIP Adv. 4, 107133. https://doi.org/10.1063/1.4899186 (2014).
4. Lepadatu, S. Boris computational spintronics—high performance multi-mesh magnetic and spin transport modeling software. J. Appl. Phys. 128, 243902 https://doi.org/10.1063/5.0024382 (2020).
5. Bruckner, F; Koraltan, S; Abert, C; Suess, D. magnum.np: a PyTorch-based GPU-enhanced finite difference micromagnetic simulation framework for high-level development and inverse design. Sci. Rep.; 2023; 13, [DOI: https://dx.doi.org/10.1038/s41598-023-39192-5] [PubMed: https://www.ncbi.nlm.nih.gov/pubmed/37491598][PubMedCentral: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC10368681]1:CAS:528:DC%2BB3sXhsFCktLrF 12054.
6. Fischbacher, T; Franchin, M; Bordignon, G; Fangohr, H. A systematic approach to multiphysics extensions of finite-element-based micromagnetic simulations: nmag. IEEE Trans. Magn.; 2007; 43, pp. 2896-2898. [DOI: https://dx.doi.org/10.1109/TMAG.2007.893843]
7. Kakay, A; Westphal, E; Hertel, R. Speedup of FEM micromagnetic simulations with graphical processing units. IEEE Trans. Magn.; 2010; 46, pp. 2303-2306. [DOI: https://dx.doi.org/10.1109/TMAG.2010.2048016]
8. Chang, R., Li, S., Lubarda, M., Livshitz, B. & Lomakin, V. FastMag: fast micromagnetic simulator for complex magnetic structures. J. Appl. Phys. 109, 07D358 (2011).
9. Bisotti, M.-A. et al. FinMag: finite-element micromagnetic simulation tool (Version 1.0). Zenodo (2018).
10. Abert, C; Exl, L; Bruckner, F; Drews, A; Suess, D. magnum.fe: a micromagnetic finite-element simulation code based on FEniCS. J. Magn. Magn. Mater.; 2013; 345, pp. 29-35. [DOI: https://dx.doi.org/10.1016/j.jmmm.2013.05.051] 1:CAS:528:DC%2BC3sXht1emt73N
11. Körber, L; Quasebarth, G; Otto, A; Kákay, A. Finite-element dynamic-matrix approach for spin-wave dispersions in magnonic waveguides with arbitrary cross section. AIP Adv.; 2021; 11, 095006. [DOI: https://dx.doi.org/10.1063/5.0054169]
12. Bruckner, F et al. Solving large-scale inverse magnetostatic problems using the adjoint method. Sci. Rep.; 2017; 7, [DOI: https://dx.doi.org/10.1038/srep40816] [PubMed: https://www.ncbi.nlm.nih.gov/pubmed/28098851][PubMedCentral: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5242220]1:CAS:528:DC%2BC2sXhsVamsL8%3D 40816.
13. Abert, C. et al. A fast finite-difference algorithm for topology optimization of permanent magnets. J. Appl. Phys. 122, 113904 (2017).
14. Huber, C. et al. Topology optimized and 3D printed polymer-bonded permanent magnets for a predefined external field. J. Appl. Phys. 122, 053904 (2017).
15. Papp, A; Porod, W; Csaba, G. Nanoscale neural network using non-linear spin-wave interference. Nat. Commun.; 2021; 12, [DOI: https://dx.doi.org/10.1038/s41467-021-26711-z] [PubMed: https://www.ncbi.nlm.nih.gov/pubmed/34741047][PubMedCentral: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC8571280]1:CAS:528:DC%2BB3MXisVamtLbI 6422.
16. Wang, Q; Chumak, AV; Pirro, P. Inverse-design magnonic devices. Nat. Commun.; 2021; 12, [DOI: https://dx.doi.org/10.1038/s41467-021-22897-4] [PubMed: https://www.ncbi.nlm.nih.gov/pubmed/33976137][PubMedCentral: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC8113576]1:CAS:528:DC%2BB3MXhtFWgt7nO 2636.
17. Yan, Z; Xing, Y; Han, X. Inverse design of magnonic filter. J. Magn. Magn. Mater.; 2022; 563, 169976. [DOI: https://dx.doi.org/10.1016/j.jmmm.2022.169976] 1:CAS:528:DC%2BB38XisF2gtr%2FO
18. Zenbaa, N. et al. A universal inverse-design magnonic device. Nat. Electron.8, 106–115 (2025).
19. Brown, WF. Micromagnetics, domains, and resonance. J. Appl. Phys.; 1959; 30, pp. S62-S69. [DOI: https://dx.doi.org/10.1063/1.2185970]
20. Schrefl, T et al. Numerical micromagnetics (Finite element method). Handb. Magn. Adv. Magn. Mater.; 2007; 2, 765–795.
21. Skomski, R. Simple Models of Magnetism. https://doi.org/10.1093/acprof:oso/9780198570752.001.0001 (Oxford University Press, 2008).
22. Abert, C. Micromagnetics and spintronics: models and numerical methods. Eur. Phys. J. B; 2019; 92, pp. 1-45. [DOI: https://dx.doi.org/10.1140/epjb/e2019-90599-6] 1:CAS:528:DC%2BC1MXhtFGqtbbM
23. Paszke, A. et al. Pytorch: an imperative style, high-performance deep learning library. In Proc. 33rd International Conference on Neural Information Processing Systems. 8026–8037 (Curran Associates Inc., 2019).
24. Bradbury, J. et al. Jax: Autograd and XLA. Astrophysics Source Code Library ascl–2111 https://ui.adsabs.harvard.edu/abs/2021ascl.soft11002B (2021).
25. Li, S et al. PyTorch distributed. Proc. VLDB Endow.; 2020; 13, pp. 3005-3018. [DOI: https://dx.doi.org/10.14778/3415478.3415530]
26. Chen, R. T., Rubanova, Y., Bettencourt, J. & Duvenaud, D. K. Neural ordinary differential equations. In: S. Bengio, H. Wallach, H. Larochelle, K. Grauman,d N. Cesa-Bianchi, R. Garnett (eds.) Advances in Neural Information Processing Systems. Vol. 31 (NeurIPS, 2018).
27. μMAG standard problem #4. https://www.ctcms.nist.gov/~rdm/std4/spec4.html (2024).
28. NeuralMag: an open-source nodal finite-difference code for inverse micromagnetics. https://gitlab.com/neuralmag/neuralmag (2024).
29. Heistracher, P; Abert, C; Bruckner, F; Schrefl, T; Suess, D. Proposal for a micromagnetic standard problem: domain wall pinning at phase boundaries. J. Magn. Magn. Mater.; 2022; 548, 168875. [DOI: https://dx.doi.org/10.1016/j.jmmm.2021.168875] 1:CAS:528:DC%2BB38XmsFSlsQ%3D%3D
30. Rozvany, G. The SIMP method in topology optimization–theoretical background, advantages and new applications. In Proc. 8th Symposium on Multidisciplinary Analysis and Optimization, Vol. 4738. https://doi.org/10.2514/6.2000-4738 (American Institute of Aeronautics and Astronautics, 2000).
31. Voronov, AA et al. Inverse-design topology optimization of magnonic devices using level-set method. npj Spintronics; 2025; 3, pp. 1-8. [DOI: https://dx.doi.org/10.1038/s44306-025-00082-3]
32. Miltat, J. E & Donahue, M. J. Numerical micromagnetics: finite difference methods. In Handbook of Magnetism and Advanced Magnetic Materials. (eds Kronmülle, H., & Parkin, S.) 742–764 (John Wiley & Sons, 2007).
33. Suess, D; Koraltan, S; Slanovc, F; Bruckner, F; Abert, C. Accurate finite-difference micromagnetics of magnets including RKKY interaction: analytical solution and comparison to standard micromagnetic codes. Phys. Rev. B; 2023; 107, 104424. [DOI: https://dx.doi.org/10.1103/PhysRevB.107.104424] 1:CAS:528:DC%2BB3sXosFOrs7k%3D
34. Berkov, DV; Ramstöcck, K; Hubert, A. Solving micromagnetic problems. towards an optimal numerical method. Phys. Status Solidi; 1993; 137, pp. 207-225. [DOI: https://dx.doi.org/10.1002/pssa.2211370118]
35. Ramstöck, K; Leibl, T; Hubert, A. Optimizing stray field computations in finite-element micromagnetics. J. Magn. Magn. Mater.; 1994; 135, pp. 97-110. [DOI: https://dx.doi.org/10.1016/0304-8853(94)90178-3]
36. Frostig, R., Johnson, M. J. & Leary, C. Compiling machine learning programs via high-level tracing. Syst. Mach. Learn.4, 9 (2018).
37. Meurer, A et al. SymPy: symbolic computing in Python. PeerJ Comput. Sci.; 2017; 3, e103. [DOI: https://dx.doi.org/10.7717/peerj-cs.103]
38. Kidger, P., Chen, R. T. & Lyons, T. J. "Hey, that’s not an ODE": faster ODE adjoints via seminorms. In: Proc.ICML, 5443–5452 (PMLR, 2021).
39. Kidger, P. On Neural Differential Equations. Ph.D. thesis (University of Oxford, 2021).
Copyright Nature Publishing Group 2025