Array API migration for Magpylib

[purepani]

Project

Magpylib

Summary

This proposal aims to implement Array API support for Magpylib. This will allow for the use of GPU computation, as well as allowing for automatic differentiation through JAX.

submitter

Satwik Pani

project lead

@OrtnerMichael

Community benefit

Magpylib provides implementations of analytical expressions for the magnetic field of permanent magnets and electrical currents, making field computation a matter of microseconds as opposed to the minutes and hours it takes to set up and run a common finite element (FE) simulation on expensive professional software like ANSYS or COMSOL. In its specific field of application, i.e. simulation of magnetostatic fields of magnets and currents, multiple studies have demonstrated the advantages of Magpylib over FE and other numerical approaches. Specifically, the fast computation times enable users to spend more time on design and variation, not having to wait for long periods between subsequent simulations. Use-cases from scientific publications include various magnetic sensor, coil, and actuator systems, Halbach array design for MRI, fusion reactor design, magnetic trap design for quantum computing, magnetic nano-particle detection, atom accelerator focus lens design and many others.

Currently, Magpylib is restricted to using NumPy, and consequently the CPU for computations. For computations on a small scale, this is quite fast, but for applications such as evaluating over a large 3D grid, or evaluating over a large number of magnet geometry dimensions, or both, the speed of this computation can still be a limiting factor(magpylib/magpylib#827, magpylib/magpylib#792, magpylib/magpylib#835).

By implementing Array API support, users can take advantage of libraries such as JAX or PyTorch to use the GPU for computations, significantly speeding up these larger computations. In addition to simulating over a large number of magnet parameters, users could also more easily simulate complex geometries, like larger triangle meshes.

In addition to the potential speed improvements, other features such as auto-differentiation in JAX can be taken advantage of for more accurate computations. Many computations in magnetism involve first or second order derivatives, and using auto-differentiation would allow for faster and more accurate derivatives compared to numerical approximations. Users would not only be able to accurately and quickly calculate the direct field, but they would also be able to directly and quickly calculate magnetic forces, and use gradient optimization much more efficiently.

Finally, due to the generic nature of the array API, we intend to upstream specialized functions(such as elliptic function implementations), either into SciPy, or upstream array libraries as needed to support this. These algorithms are usable well beyond this library, and users of SciPy and the upstream libraries should benefit from this work as well.

Amount requested (USD)

10000

Execution plan

The plan is intended to be done by @purepani and @jack-smithy, with an intended 75/25 split on work for the budget.

The library is currently split into a core field computation layer, a functional interface, and an object oriented interface, with the core field interface supporting the functional and object oriented interface; as such, we will first support the Array API in the core field computation layer before expanding to the rest of the library.

We aim to support at least NumPy, PyTorch, and JAX array inputs. In particular, we would like to support JAX jit compilation and auto-differentiation, which will open a wide range of new opportunities for users. Supporting PyTorch auto-differentiation and compilation is out of scope for this grant, due to the seemingly higher complexity, but is intended as a future step.

This plan will be split into the following four steps:

1. Improvement/Expansion of testing of core field computation
   For testing, the lowest level, the core field computation, lacks sufficient testing for a conversion to the Array API. Tests at this level will need to be added here before converting to the Array API. These tests have already been needed from core function tests based on FE magpylib/magpylib#732.

   Budget: 20h / $1000
   Time: 1 month

2. Core field computation conversion.
   This step will involve converting all the tests for the core field computation level to support testing the Array API, before proceeding to convert all the core fields to support the Array API. Some of the work of this section has already been done, and is functional for the Array API, but certain algorithms, particularly the elliptic integrals, are difficult implement generically, and require some experimentation on convergence before implementation. The algorithm implementations and any generic tooling can be upstreamed to the SciPy or the upstream array libraries as necessary.

   Budget: 100h/$5000
   TIme: 6 months

3. Functional Interface and Object Oriented Interface Conversion
   At this level, the current interface uses the SciPy Rotation class as part of the interface for representing rotations of both magnetic sources and magnetic observers. In nightly SciPy, the Rotation class has been converted to support the Array API, slated for the next major version of SciPy. This step will therefore involve creating infrastructure to test with nightly SciPy, in addition to adding Array API support to the tests and the code to support the Array API.

   Budget: 40h/$2000
   Time: 3 months

4. Documentation
   In addition to making sure to clearly document library support, we hope to add informative documentation on how to use the Array API in the context of Magpylib. The library is intended to be used by scientists first and foremost, so we hope to write clear tutorials and examples, and make sure that the users will be made easily aware of the Array API and it's advantages. The examples will, in particular, aim to showcase the use of JAX features in the context of magnetic field computations with auto-differentiation and GPU computations. The examples will demonstrate speed improvements and using auto-differentiation for simple force computation.

   Budget: 40h/$2000
   Time: 2 months

[OrtnerMichael]

I confirm Magpylib's strong support for @purepani 's proposal!

br,
michael ortner