Magnetix

Magnetic fields of current-carrying wires — differentiable, GPU-ready, in JAX.

Magnetix computes the magnetic field B and vector potential A produced by current-carrying filaments, at arbitrary points in space. Everything is a plain JAX function, so derivatives with respect to conductor geometry or current follow from jax.grad and are exact.

It is a building block. Objectives, optimisers, and derived engineering quantities such as force, torque and inductance are left to the caller.

Example

import jax.numpy as jnp
import magnetix as mx

angle = jnp.linspace(0, 2*jnp.pi, 201)
loop  = jnp.stack([0.05*jnp.cos(angle), 0.05*jnp.sin(angle), jnp.zeros_like(angle)], -1)
coil  = mx.polyline(loop, current=400.0)          # 5 cm radius, 400 A

z       = jnp.linspace(-0.05, 0.05, 400)
targets = jnp.stack([jnp.zeros_like(z), jnp.zeros_like(z), z], -1)

B = mx.B(coil, targets)     # (400, 3) tesla
A = mx.A(coil, targets)     # (400, 3) tesla-metre

Derivatives propagate through the code that builds the coil:

import jax

def central_field(radius):
    angle = jnp.linspace(0, 2*jnp.pi, 201)
    loop = jnp.stack([radius*jnp.cos(angle), radius*jnp.sin(angle), jnp.zeros_like(angle)], -1)
    return mx.B(mx.polyline(loop, 400.0), jnp.zeros((1, 3)))[0, 2]

jax.grad(central_field)(0.05)      # -0.1005 T/m, exact

A longer walkthrough with plots is in examples/demo.ipynb.

Install

pip install magnetix

The only dependencies are jax and jaxtyping. The same code runs on CPU, GPU and TPU, so it is enough to install the jax wheel that matches the hardware.

JAX computes in single precision by default, which resolves ordinary geometries to about six digits. Precision degrades close to a conductor, where the field emerges from a cancellation between large terms. The default softening length holds this in check, so it becomes relevant for filaments given a much smaller eps. Double precision removes it, and has to be enabled before the first array is created:

import jax
jax.config.update("jax_enable_x64", True)

Scope

Sources are straight current segments, softened so that the field and its gradient stay finite on the conductor. The field is a superposition over segments, each integrated in closed form,

$$ \mathbf{B}(\mathbf{r}) = \frac{\mu_0}{4\pi} \sum_i I_i , \frac{R_{a,i} + R_{b,i}} {R_{a,i} R_{b,i} \left( R_{a,i} R_{b,i} + \mathbf{a}_i \cdot \mathbf{b}_i \right)} , \mathbf{a}_i \times \mathbf{b}_i , $$

where segment $i$ runs from $\mathbf{r}{1,i}$ to $\mathbf{r}{2,i}$ carrying current $I_i$, and $\mathbf{a}i = \mathbf{r} - \mathbf{r}{1,i}$, $\mathbf{b}i = \mathbf{r} - \mathbf{r}{2,i}$, $R_{a,i} = \sqrt{\lVert \mathbf{a}i \rVert^2 + \varepsilon_i^2}$ and $R{b,i} = \sqrt{\lVert \mathbf{b}_i \rVert^2 + \varepsilon_i^2}$ with softening length $\varepsilon_i$.

polyline builds segments from any curve, which covers coils, racetracks, saddle windings and helices. All functions compose with jit, grad, vmap and jacfwd. Geometries are air-core, and the library evaluates the field of a prescribed current distribution.

The integral along each segment is exact, so accuracy is governed by how well the filaments represent the conductor. Beyond roughly one conductor width this is straightforward. Closer in, and inside a winding pack, a bundle of filaments converges slowly and the softening length begins to dominate the answer, which leaves the current version unsuited to peak-field and force calculations within the conductor itself. Cross-section quadrature and a dedicated near-conductor treatment are the natural extensions.

Summation is direct, so cost grows with the product of source and target counts, while tiling keeps memory bounded. Version 0.x is small and the API may still change; the kernels are validated against closed-form solutions.

License

MIT.