helicon.fields

Magnetic field computation: Biot-Savart solver, field line tracing, FRC topology, and caching.

Core API

`helicon.fields.biot_savart.Coil(z, r, I)` `dataclass`

A single circular current loop in SI units.

Parameters:

Name	Type	Description	Default
`z`	`float`	Axial position of the coil centre [m].	required
`r`	`float`	Radius of the coil [m].	required
`I`	`float`	Current (ampere-turns) [A].	required

`helicon.fields.biot_savart.Grid(z_min, z_max, r_max, nz, nr)` `dataclass`

Axisymmetric (r, z) computation grid.

The grid spans z ∈ [z_min, z_max] and r ∈ [0, r_max].

Parameters:

Name	Type	Description	Default
`z_min`	`float`	Axial extent [m].	required
`z_max`	`float`	Axial extent [m].	required
`r_max`	`float`	Radial extent [m].	required
`nz`	`int`	Number of grid points along each axis.	required
`nr`	`int`	Number of grid points along each axis.	required

`helicon.fields.biot_savart.BField(Br, Bz, r, z, coils, backend)` `dataclass`

Result container — always stores NumPy arrays for interop.

Attributes:

Name	Type	Description
`Br, Bz`	`ndarray`	Radial and axial magnetic field components [T], shape `(nr, nz)`.
`r, z`	`ndarray`	1-D coordinate arrays [m].
`coils`	`list[Coil]`	Coil configuration that produced this field.
`backend`	`str`	Backend used (`"numpy"` or `"mlx"`).

Functions

`load(path)` `classmethod`

Load a previously saved BField from HDF5.

Source code in src/helicon/fields/biot_savart.py

@classmethod
def load(cls, path: str) -> BField:
    """Load a previously saved BField from HDF5."""
    if not HAS_H5PY:
        raise ImportError("h5py is required for BField.load()")
    with h5py.File(path, "r") as f:
        Br = f["Br"][:]
        Bz = f["Bz"][:]
        r = f["r"][:]
        z = f["z"][:]
        backend = f.attrs["backend"]
        n_coils = int(f.attrs["n_coils"])
        coils = [
            Coil(
                z=float(f.attrs[f"coil_{i}_z"]),
                r=float(f.attrs[f"coil_{i}_r"]),
                I=float(f.attrs[f"coil_{i}_I"]),
            )
            for i in range(n_coils)
        ]
    return cls(Br=Br, Bz=Bz, r=r, z=z, coils=coils, backend=backend)

`plot(*, component='Bz', field_lines=True, n_field_lines=8, ax=None, cmap='RdBu_r', figsize=(10, 5))`

Plot the B-field on the (z, r) plane.

Parameters:

Name	Type	Description	Default
`component`	``"Bz"`` \| ``"Br"`` \| ``"Bmag"``	Field component to show as a filled contour.	`'Bz'`
`field_lines`	`bool`	Overlay field lines via contours of the flux function ψ.	`True`
`n_field_lines`	`int`	Number of field line contours to draw.	`8`
`ax`	`matplotlib Axes`	Axes to draw on. Creates new figure if None.	`None`
`cmap`	`str`	Matplotlib colormap.	`'RdBu_r'`
`figsize`	`tuple`	Figure size when creating a new figure.	`(10, 5)`

Returns:

Type	Description
`(fig, ax)`	Matplotlib figure and axes.

Source code in src/helicon/fields/biot_savart.py

def plot(
    self,
    *,
    component: str = "Bz",
    field_lines: bool = True,
    n_field_lines: int = 8,
    ax=None,
    cmap: str = "RdBu_r",
    figsize: tuple = (10, 5),
):
    """Plot the B-field on the (z, r) plane.

    Parameters
    ----------
    component : ``"Bz"`` | ``"Br"`` | ``"Bmag"``
        Field component to show as a filled contour.
    field_lines : bool
        Overlay field lines via contours of the flux function ψ.
    n_field_lines : int
        Number of field line contours to draw.
    ax : matplotlib Axes, optional
        Axes to draw on. Creates new figure if None.
    cmap : str
        Matplotlib colormap.
    figsize : tuple
        Figure size when creating a new figure.

    Returns
    -------
    fig, ax
        Matplotlib figure and axes.
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise ImportError("matplotlib is required for BField.plot()") from exc

    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = ax.get_figure()

    Z, R = np.meshgrid(self.z, self.r)

    if component == "Bz":
        values = self.Bz
        label = "$B_z$ (T)"
    elif component == "Br":
        values = self.Br
        label = "$B_r$ (T)"
    elif component == "Bmag":
        values = np.sqrt(self.Br**2 + self.Bz**2)
        label = "$|B|$ (T)"
    else:
        raise ValueError(f"Unknown component {component!r}. Use 'Bz', 'Br', or 'Bmag'.")

    pcm = ax.pcolormesh(Z, R, values, cmap=cmap, shading="auto")
    fig.colorbar(pcm, ax=ax, label=label)

    if field_lines:
        # Flux function ψ(r,z) = cumulative integral of r*B_z over r (up to sign)
        dr = self.r[1] - self.r[0] if len(self.r) > 1 else 1.0
        psi = np.cumsum(self.r[:, None] * self.Bz, axis=0) * dr
        with contextlib.suppress(Exception):
            ax.contour(Z, R, psi, n_field_lines, colors="k", linewidths=0.6, alpha=0.5)

    # Mark coil positions
    for coil in self.coils:
        ax.plot(coil.z, coil.r, "ko", ms=6, zorder=5)
        ax.plot(coil.z, -coil.r if min(self.r) < 0 else coil.r, "ko", ms=6, zorder=5)

    ax.set_xlabel("z (m)")
    ax.set_ylabel("r (m)")
    ax.set_title(f"Magnetic Field — {component}")

    return fig, ax

`save(path)`

Write field data + coil metadata to an HDF5 file.

Source code in src/helicon/fields/biot_savart.py

def save(self, path: str) -> None:
    """Write field data + coil metadata to an HDF5 file."""
    if not HAS_H5PY:
        raise ImportError("h5py is required for BField.save()")
    with h5py.File(path, "w") as f:
        f.create_dataset("Br", data=self.Br)
        f.create_dataset("Bz", data=self.Bz)
        f.create_dataset("r", data=self.r)
        f.create_dataset("z", data=self.z)
        f.attrs["backend"] = self.backend
        f.attrs["n_coils"] = len(self.coils)
        for i, c in enumerate(self.coils):
            f.attrs[f"coil_{i}_z"] = c.z
            f.attrs[f"coil_{i}_r"] = c.r
            f.attrs[f"coil_{i}_I"] = c.I

`helicon.fields.biot_savart.compute_bfield(coils, grid, *, backend='auto', n_phi=128, differentiable=False)`

Compute the applied magnetic field from a set of circular coils.

Parameters:

Name	Type	Description	Default
`coils`	`sequence of Coil`	Coil definitions in SI units.	required
`grid`	`Grid`	Computational domain.	required
`backend`	``"auto"`` \| ``"mlx"`` \| ``"numpy"``	`"auto"` uses MLX if available, else NumPy.	`'auto'`
`n_phi`	`int`	Azimuthal quadrature points (MLX backend only).	`128`
`differentiable`	`bool`	When True, use the loop-based MLX path so the computation graph is preserved for `mx.grad()`. When False (default), use the vectorized + compiled path for faster non-gradient evaluation.	`False`

Returns:

Type	Description
`BField`	Result container with NumPy arrays.

Source code in src/helicon/fields/biot_savart.py

def compute_bfield(
    coils: Sequence[Coil],
    grid: Grid,
    *,
    backend: str = "auto",
    n_phi: int = 128,
    differentiable: bool = False,
) -> BField:
    """Compute the applied magnetic field from a set of circular coils.

    Parameters
    ----------
    coils : sequence of Coil
        Coil definitions in SI units.
    grid : Grid
        Computational domain.
    backend : ``"auto"`` | ``"mlx"`` | ``"numpy"``
        ``"auto"`` uses MLX if available, else NumPy.
    n_phi : int
        Azimuthal quadrature points (MLX backend only).
    differentiable : bool
        When True, use the loop-based MLX path so the computation graph is
        preserved for ``mx.grad()``.  When False (default), use the
        vectorized + compiled path for faster non-gradient evaluation.

    Returns
    -------
    BField
        Result container with NumPy arrays.
    """
    coils = list(coils)

    if backend == "auto":
        backend = "mlx" if HAS_MLX else "numpy"

    if backend == "mlx":
        if not HAS_MLX:
            raise ImportError("MLX is not installed. Install with: pip install 'helicon[mlx]'")
        Br, Bz, r, z = _compute_mlx(coils, grid, n_phi=n_phi, differentiable=differentiable)
    elif backend == "numpy":
        Br, Bz, r, z = _compute_numpy(coils, grid)
    else:
        raise ValueError(f"Unknown backend: {backend!r}. Use 'auto', 'mlx', or 'numpy'.")

    return BField(Br=Br, Bz=Bz, r=r, z=z, coils=coils, backend=backend)

`helicon.fields.biot_savart.compute_bfield_mlx_differentiable(coil_params, grid_r, grid_z, n_phi=128)`

MLX-native differentiable Biot-Savart computation.

Designed for composition with mlx.core.grad(). Does not call mx.eval so the computation graph is preserved.

Parameters:

Name	Type	Description	Default
`coil_params`	`(array, shape(N_coils, 3))`	Each row is `[z_coil, radius, current]` in SI.	required
`grid_r`	`(array, shape(N_pts))`	Flattened grid coordinates.	required
`grid_z`	`(array, shape(N_pts))`	Flattened grid coordinates.	required
`n_phi`	`int`	Number of azimuthal quadrature points.	`128`

Returns:

Type	Description
`Br, Bz : mx.array, shape (N_pts,)`

Source code in src/helicon/fields/biot_savart.py

def compute_bfield_mlx_differentiable(
    coil_params: mx.array,
    grid_r: mx.array,
    grid_z: mx.array,
    n_phi: int = 128,
) -> tuple[mx.array, mx.array]:
    """MLX-native differentiable Biot-Savart computation.

    Designed for composition with ``mlx.core.grad()``.  Does **not** call
    ``mx.eval`` so the computation graph is preserved.

    Parameters
    ----------
    coil_params : mx.array, shape (N_coils, 3)
        Each row is ``[z_coil, radius, current]`` in SI.
    grid_r, grid_z : mx.array, shape (N_pts,)
        Flattened grid coordinates.
    n_phi : int
        Number of azimuthal quadrature points.

    Returns
    -------
    Br, Bz : mx.array, shape (N_pts,)
    """
    if not HAS_MLX:
        raise ImportError("MLX is required for compute_bfield_mlx_differentiable")

    phi = mx.linspace(0.0, 2.0 * math.pi, n_phi + 1)[:-1]  # exclude endpoint
    cos_phi = mx.cos(phi)  # (N_phi,)

    N_pts = grid_r.shape[0]
    N_coils = coil_params.shape[0]

    Br_total = mx.zeros((N_pts,))
    Bz_total = mx.zeros((N_pts,))

    for i in range(N_coils):
        z_c = coil_params[i, 0]
        a = coil_params[i, 1]
        I_c = coil_params[i, 2]

        dz = grid_z - z_c  # (N_pts,)

        # Broadcast: r -> (N_pts, 1), cos_phi -> (1, N_phi)
        r_2d = mx.reshape(grid_r, (N_pts, 1))
        dz_2d = mx.reshape(dz, (N_pts, 1))
        cos_2d = mx.reshape(cos_phi, (1, n_phi))

        R_sq = r_2d * r_2d + a * a - 2.0 * a * r_2d * cos_2d + dz_2d * dz_2d
        R_sq = mx.maximum(R_sq, 1e-20)  # singularity guard

        R_inv3 = R_sq ** (-1.5)

        prefactor = MU_0 * I_c * a / (2.0 * n_phi)

        # Sum over phi (axis=1)
        Br_total = Br_total + prefactor * mx.sum(cos_2d * dz_2d * R_inv3, axis=1)
        Bz_total = Bz_total + prefactor * mx.sum((a - r_2d * cos_2d) * R_inv3, axis=1)

    return Br_total, Bz_total