Adding multipole magnet utilities

beamphysics/fields/__init__.py

@@ -1,5 +1,7 @@
 from .fieldmesh import FieldMesh
+from . import multipole
 
 __all__ = [
     "FieldMesh",
+    "multipole",
 ]

beamphysics/fields/multipole.py

@@ -0,0 +1,179 @@
+from functools import partial
+from math import factorial
+
+import numpy as np
+from scipy.integrate import quad
+from scipy.optimize import curve_fit
+
+
+def synthesize_field(
+    x: np.ndarray | float,
+    y: np.ndarray | float,
+    multipoles: list[tuple[float, float]],
+) -> tuple[np.ndarray | float, np.ndarray | float]:
+    r"""
+    Calculate transverse magnetic field components at position (x, y)
+    from multipole coefficients.
+
+    The magnetic field is computed from the complex multipole expansion:
+
+    .. math::
+        B_y + i B_x = \sum_{n=0}^{N} \frac{C_n}{n!} z^n
+
+    where :math:`z = x + iy`, :math:`C_n = B_n + i S_n`, and the sum runs
+    over the provided multipole coefficients.
+
+    Parameters
+    ----------
+    x : float or np.ndarray
+        Horizontal position in meters.
+    y : float or np.ndarray
+        Vertical position in meters.
+    multipoles : list of tuple[float, float]
+        Multipole coefficients as (B_n, S_n) pairs, where
+        n=0 is dipole, n=1 is quadrupole, n=2 is sextupole, etc.
+        B_n is the normal component and S_n is the skew component,
+        both in units of T/m^n.
+
+    Returns
+    -------
+    B_x : float or np.ndarray
+        Horizontal magnetic field component in Tesla.
+    B_y : float or np.ndarray
+        Vertical magnetic field component in Tesla.
+    """
+    z = x + 1j * y
+
+    B_complex = 0j
+    for n, (B_n, S_n) in enumerate(multipoles):
+        C_n = B_n + 1j * S_n
+        B_complex = B_complex + (C_n / factorial(n)) * z**n
+
+    B_y = np.real(B_complex)
+    B_x = np.imag(B_complex)
+
+    return B_x, B_y
+
+
+def decompose_field(
+    data: list[tuple[float, float]] | np.ndarray,
+    r0: float,
+    nmax: int,
+) -> list[tuple[float, float]]:
+    r"""
+    Decompose azimuthal field measurements on a circle into multipole
+    coefficients using a least-squares fit.
+
+    Given measurements of the tangential field component
+    :math:`B_\phi(\phi)` at radius :math:`r_0`, this function fits
+    the model:
+
+    .. math::
+        B_\phi(\phi) = \sum_{n=0}^{n_\mathrm{max}} \frac{r_0^n}{n!}
+        \left[ B_n \cos\!\left((n+1)\phi\right)
+        - S_n \sin\!\left((n+1)\phi\right) \right]
+
+    Parameters
+    ----------
+    data : list of tuple[float, float] or np.ndarray
+        Measurement data as (phi, B_phi) pairs, where phi is
+        the azimuthal angle in radians and B_phi is the tangential
+        field component in Tesla.
+    r0 : float
+        Measurement radius in meters.
+    nmax : int
+        Maximum multipole order to fit. 0 = dipole, 1 = quadrupole,
+        2 = sextupole, etc.
+
+    Returns
+    -------
+    list of tuple[float, float]
+        Multipole coefficients as (B_n, S_n) pairs for n = 0 to nmax.
+        B_n is the normal component and S_n is the skew component,
+        both in units of T/m^n.
+    """
+    data_array = np.array(data)
+    phi = data_array[:, 0]
+    B_phi_measured = data_array[:, 1]
+
+    def model(phi_vals, *params):
+        B_phi = np.zeros_like(phi_vals)
+        for n in range(nmax + 1):
+            B_n = params[2 * n]
+            S_n = params[2 * n + 1]
+            factor = (r0**n) / factorial(n)
+            B_phi += factor * (
+                B_n * np.cos((n + 1) * phi_vals)
+                - S_n * np.sin((n + 1) * phi_vals)
+            )
+        return B_phi
+
+    initial_params = np.zeros(2 * (nmax + 1))
+    popt, _ = curve_fit(model, phi, B_phi_measured, p0=initial_params)
+
+    multipoles = []
+    for n in range(nmax + 1):
+        B_n = popt[2 * n]
+        S_n = popt[2 * n + 1]
+        multipoles.append((B_n, S_n))
+
+    return multipoles
+
+
+def _integrand(
+    phi: float,
+    B_actual: callable,
+    B_design: callable,
+    r0: float,
+) -> float:
+    """Squared relative field error at azimuthal angle phi."""
+    x = r0 * np.cos(phi)
+    y = r0 * np.sin(phi)
+
+    B_x_actual, B_y_actual = B_actual(x, y)
+    B_x_design, B_y_design = B_design(x, y)
+
+    B_design_mag_sq = B_x_design**2 + B_y_design**2
+    error_mag_sq = (B_x_actual - B_x_design) ** 2 + (B_y_actual - B_y_design) ** 2
+
+    return error_mag_sq / B_design_mag_sq
+
+
+def scalar_error(
+    B_actual: callable,
+    B_design: callable,
+    r0: float,
+) -> float:
+    r"""
+    Compute the normalized RMS field error between two field
+    distributions on a circle.
+
+    Evaluates:
+
+    .. math::
+        \Delta B / B = \sqrt{
+            \frac{1}{2\pi} \int_0^{2\pi}
+            \frac{|\mathbf{B}_\mathrm{actual} - \mathbf{B}_\mathrm{design}|^2}
+            {|\mathbf{B}_\mathrm{design}|^2}
+            \, d\phi
+        }
+
+    Parameters
+    ----------
+    B_actual : callable
+        Field function with signature ``(x, y) -> (B_x, B_y)``
+        returning the actual field in Tesla.
+    B_design : callable
+        Field function with signature ``(x, y) -> (B_x, B_y)``
+        returning the design field in Tesla.
+    r0 : float
+        Evaluation radius in meters.
+
+    Returns
+    -------
+    float
+        Normalized RMS field error (dimensionless).
+    """
+    integrand = partial(_integrand, B_actual=B_actual, B_design=B_design, r0=r0)
+    integral_result, _ = quad(integrand, 0, 2 * np.pi)
+    return np.sqrt(integral_result / (2 * np.pi))

docs/api/fields.md

@@ -1 +1,3 @@
 ::: beamphysics.FieldMesh
+
+::: beamphysics.fields.multipole

mkdocs.yml

@@ -20,6 +20,7 @@ nav:
           - examples/fields/field_conversion.ipynb
           - examples/fields/corrector_modeling.ipynb
           - examples/fields/solenoid_modeling.ipynb
+          - examples/fields/multipole_utils.ipynb
       - Wakefields:
           - examples/wakefields/resistive_wall.ipynb
           - examples/wakefields/impedance_wakefield.ipynb

tests/test_multipole.py

@@ -0,0 +1,87 @@
+import numpy as np
+import pytest
+from functools import partial
+
+from beamphysics.fields.multipole import decompose_field, scalar_error, synthesize_field
+
+
+def test_synthesize_pure_quadrupole():
+    """A pure normal quadrupole B_1=1 T/m should give B_y = x, B_x = y."""
+    multipoles = [(0.0, 0.0), (1.0, 0.0)]
+    x = np.linspace(-0.02, 0.02, 10)
+    y = np.linspace(-0.02, 0.02, 10)
+    X, Y = np.meshgrid(x, y)
+
+    B_x, B_y = synthesize_field(X, Y, multipoles)
+
+    np.testing.assert_allclose(B_y, X, atol=1e-15)
+    np.testing.assert_allclose(B_x, Y, atol=1e-15)
+
+
+def test_synthesize_pure_dipole():
+    """A pure normal dipole B_0=1 T should give uniform B_y=1, B_x=0."""
+    multipoles = [(1.0, 0.0)]
+    B_x, B_y = synthesize_field(0.01, 0.005, multipoles)
+
+    np.testing.assert_allclose(B_y, 1.0, atol=1e-15)
+    np.testing.assert_allclose(B_x, 0.0, atol=1e-15)
+
+
+def test_synthesize_skew_dipole():
+    """A pure skew dipole S_0=1 T should give uniform B_x=1, B_y=0."""
+    multipoles = [(0.0, 1.0)]
+    B_x, B_y = synthesize_field(0.01, 0.005, multipoles)
+
+    np.testing.assert_allclose(B_x, 1.0, atol=1e-15)
+    np.testing.assert_allclose(B_y, 0.0, atol=1e-15)
+
+
+def test_decompose_roundtrip():
+    """Decompose should recover the multipoles used to synthesize B_phi."""
+    ground_truth = [
+        (1.0, 0.5),
+        (10.0, -5.0),
+        (100.0, 60.0),
+    ]
+    r0 = 0.01
+    phi_samples = np.linspace(0, 2 * np.pi, 72, endpoint=False)
+
+    data = []
+    for phi in phi_samples:
+        x = r0 * np.cos(phi)
+        y = r0 * np.sin(phi)
+        Bx, By = synthesize_field(x, y, ground_truth)
+        B_phi = -Bx * np.sin(phi) + By * np.cos(phi)
+        data.append((phi, B_phi))
+
+    recovered = decompose_field(data, r0, nmax=2)
+
+    for n, (Bn_true, Sn_true) in enumerate(ground_truth):
+        Bn_rec, Sn_rec = recovered[n]
+        np.testing.assert_allclose(Bn_rec, Bn_true, atol=1e-10)
+        np.testing.assert_allclose(Sn_rec, Sn_true, atol=1e-10)
+
+
+def test_scalar_error_identical_fields():
+    """Identical fields should give zero error."""
+    multipoles = [(0.0, 0.0), (10.0, 0.0)]
+    B = partial(synthesize_field, multipoles=multipoles)
+
+    error = scalar_error(B, B, r0=0.01)
+    np.testing.assert_allclose(error, 0.0, atol=1e-14)
+
+
+def test_scalar_error_with_octupole():
+    """Reproduce the notebook's octupole error example."""
+    design_multipoles = [(0.0, 0.0), (10.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
+    error_multipoles = [(0.0, 0.0), (0.0, 0.0), (0.0, 0.0), (80.0, 0.0)]
+    actual_multipoles = [
+        (d[0] + e[0], d[1] + e[1])
+        for d, e in zip(design_multipoles, error_multipoles)
+    ]
+
+    B_design = partial(synthesize_field, multipoles=design_multipoles)
+    B_actual = partial(synthesize_field, multipoles=actual_multipoles)
+
+    error = scalar_error(B_actual, B_design, r0=0.010)
+    np.testing.assert_allclose(error, 1 / 7500, rtol=1e-10)