Adds array api support for circle BHMJ layer

magpylib · purepani · May 20, 2025 · May 6, 2025 · May 6, 2025 · May 19, 2025
commit bd375a12e93f63da1e209cbaa46859bb4e0031e6
diff --git a/src/magpylib/_src/fields/field_BH_circle.py b/src/magpylib/_src/fields/field_BH_circle.py
@@ -6,9 +6,8 @@
 from __future__ import annotations
 
 import numpy as np
-from scipy.constants import mu_0 as MU0
-
 from array_api_compat import array_namespace
+from scipy.constants import mu_0 as MU0
 
 from magpylib._src.fields.special_cel import cel_iter
 from magpylib._src.input_checks import check_field_input
@@ -114,34 +113,39 @@ def BHJM_circle(
     - treat special cases
     """
 
+    xp = array_namespace(observers, diameter, current)
+    observers = xp.astype(observers, xp.float64)
+    diameter = xp.astype(diameter, xp.float64)
+    current = xp.astype(current, xp.float64)
+
     # allocate
-    BHJM = np.zeros_like(observers, dtype=float)
+    BHJM = xp.zeros_like(observers, dtype=xp.float64)
 
     check_field_input(field)
     if field in "MJ":
         return BHJM
 
     r, phi, z = cart_to_cyl_coordinates(observers)
-    r0 = np.abs(diameter / 2)
+    r0 = xp.abs(diameter / 2)
 
     # Special cases:
     # case1: loop radius is 0 -> return (0,0,0)
     mask1 = r0 == 0
     # case2: at singularity -> return (0,0,0)
-    mask2 = np.logical_and(abs(r - r0) < 1e-15 * r0, z == 0)
+    mask2 = xp.logical_and(xp.abs(r - r0) < 1e-15 * r0, z == 0)
     # case3: r=0
     mask3 = r == 0
-    if np.any(mask3):
-        mask4 = mask3 * ~mask1  # only relevant if not also case1
-        BHJM[mask4, 2] = (
+    if xp.any(mask3):
+        mask4 = mask3 & ~mask1  # only relevant if not also case1
+        BHJM[:, 2][mask4] = (
             (r0[mask4] ** 2 / (z[mask4] ** 2 + r0[mask4] ** 2) ** (3 / 2))
             * current[mask4]
             * 0.5
         )
 
     # general case
-    mask5 = ~np.logical_or(np.logical_or(mask1, mask2), mask3)
-    if np.any(mask5):
+    mask5 = ~xp.logical_or(xp.logical_or(mask1, mask2), mask3)
+    if xp.any(mask5):
         BHJM[mask5] = current_circle_Hfield(
             r0=r0[mask5],
             r=r[mask5],

diff --git a/src/magpylib/_src/utility.py b/src/magpylib/_src/utility.py
@@ -13,6 +13,7 @@
 from math import log10
 
 import numpy as np
+from array_api_compat import array_namespace
 
 from magpylib._src.exceptions import MagpylibBadUserInput
 
@@ -328,21 +329,23 @@ def cart_to_cyl_coordinates(observer):
     cartesian observer positions to cylindrical coordinates
     observer: ndarray, shape (n,3)
     """
-    x, y, z = observer.T
-    r, phi = np.sqrt(x**2 + y**2), np.arctan2(y, x)
+    xp = array_namespace(observer)
+    x, y, z = observer[..., 0], observer[..., 1], observer[..., 2]
+    r, phi = xp.sqrt(x**2 + y**2), xp.atan2(y, x)
     return r, phi, z
 
 
 def cyl_field_to_cart(phi, Br, Bphi=None):
     """
     transform Br,Bphi to Bx, By
     """
+    xp = array_namespace(phi, Br, Bphi)
     if Bphi is not None:
-        Bx = Br * np.cos(phi) - Bphi * np.sin(phi)
-        By = Br * np.sin(phi) + Bphi * np.cos(phi)
+        Bx = Br * xp.cos(phi) - Bphi * xp.sin(phi)
+        By = Br * xp.sin(phi) + Bphi * xp.cos(phi)
     else:
-        Bx = Br * np.cos(phi)
-        By = Br * np.sin(phi)
+        Bx = Br * xp.cos(phi)
+        By = Br * xp.sin(phi)
 
     return Bx, By
 

diff --git a/tests/test_BHMJ_level.py b/tests/test_BHMJ_level.py
@@ -453,14 +453,14 @@ def test_BHJM_magnet_trimesh_BH():
 def test_BHJM_circle():
     """Test of current circle field core function"""
     kw = {
-        "observers": np.array([(1, 1, 1), (2, 2, 2), (3, 3, 3)]),
-        "current": np.array([1, 1, 2]) * 1e3,
-        "diameter": np.array([2, 4, 6]),
+        "observers": xp.asarray([(1, 1, 1), (2, 2, 2), (3, 3, 3)]),
+        "current": xp.asarray([1.0, 1.0, 2.0]) * 1e3,
+        "diameter": xp.asarray([2, 4, 6]),
     }
     H, B, M, _ = helper_check_HBMJ_consistency(BHJM_circle, **kw)
 
     Btest = (
-        np.array(
+        xp.asarray(
             [
                 [0.06235974, 0.06235974, 0.02669778],
                 [0.03117987, 0.03117987, 0.01334889],
@@ -472,7 +472,7 @@ def test_BHJM_circle():
     np.testing.assert_allclose(B, Btest)
 
     Htest = (
-        np.array(
+        xp.asarray(
             [
                 [49624.3033947, 49624.3033947, 21245.41908818],
                 [24812.15169735, 24812.15169735, 10622.70954409],
@@ -564,9 +564,9 @@ def test_BHJM_dipole():
 
 def test_field_loop_specials():
     """test loop special cases"""
-    cur = np.array([1, 1, 1, 1, 0, 2])
-    dia = np.array([2, 2, 0, 0, 2, 2])
-    obs = np.array([(0, 0, 0), (1, 0, 0), (0, 0, 0), (1, 0, 0), (1, 0, 0), (0, 0, 0)])
+    cur = xp.asarray([1, 1, 1, 1, 0, 2])
+    dia = xp.asarray([2, 2, 0, 0, 2, 2])
+    obs = xp.asarray([(0, 0, 0), (1, 0, 0), (0, 0, 0), (1, 0, 0), (1, 0, 0), (0, 0, 0)])
 
     B = (
         BHJM_circle(
@@ -708,28 +708,28 @@ def test_field_line_special_cases():
 
 def test_field_loop2():
     """test if field function accepts correct inputs"""
-    curr = np.array([1])
-    dia = np.array([2])
-    obs = np.array([[0, 0, 0]])
+    curr = xp.asarray([1])
+    dia = xp.asarray([2])
+    obs = xp.asarray([[0, 0, 0]])
     B = BHJM_circle(
         field="B",
         observers=obs,
         current=curr,
         diameter=dia,
     )
 
-    curr = np.array([1] * 2)
-    dia = np.array([2] * 2)
-    obs = np.array([[0, 0, 0]] * 2)
+    curr = xp.asarray([1] * 2)
+    dia = xp.asarray([2] * 2)
+    obs = xp.asarray([[0, 0, 0]] * 2)
     B2 = BHJM_circle(
         field="B",
         observers=obs,
         current=curr,
         diameter=dia,
     )
 
-    assert_allclose(B, (B2[0],))
-    assert_allclose(B, (B2[1],))
+    assert_allclose(B, (B2[0, ...],))
+    assert_allclose(B, (B2[1, ...],))
 
 
 def test_field_line_from_vert():