magnum.np -- A PyTorch based GPU enhanced Finite Difference Micromagnetic Simulation Framework for High Level Development and Inverse Design

Abstract

magnum.np is a micromagnetic finite-difference library completely based on the tensor library PyTorch. The use of such a high level library leads to a highly maintainable and extensible code base which is the ideal candidate for the investigation of novel algorithms and modeling approaches. On the other hand magnum.np benefits from the devices abstraction and optimizations of PyTorch enabling the efficient execution of micromagnetic simulations on a number of computational platforms including GPU and potentially TPU systems. We demonstrate a competitive performance to state-of-the art micromagnetic codes such a mumax3 and show how our code enables the rapid implementation of new functionality. Furthermore, handling inverse problems becomes possible by using PyTorch's autograd feature.