Improves array api support for sphere core function

purepani · purepani · commit e68e9a04a192 · 2025-06-01T19:24:34.000-05:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -39,6 +39,7 @@ dependencies = [
     "plotly>=5.16",
     "array-api-compat>=1.11.2",
     "array-api-extra>=0.7.1",
+    "typing-extensions>=4.13.2",
 ]
 
 [dependency-groups]
@@ -71,6 +72,9 @@ test = [
   "jupyterlab",
   "anywidget",
   "array-api-strict>=2.3.1",
+  "jax",
+  "numpy",
+  "dask[array]",
 ]
 binder = [
     "jupytext",
diff --git a/src/magpylib/_src/fields/field_BH_sphere.py b/src/magpylib/_src/fields/field_BH_sphere.py
@@ -10,8 +10,8 @@
 from array_api_compat import array_namespace
 from scipy.constants import mu_0 as MU0
 
-from magpylib._src.input_checks import check_field_input
 from magpylib._src.array_api_utils import xp_promote
+from magpylib._src.input_checks import check_field_input
 
 
 # CORE
@@ -62,7 +62,9 @@ def magnet_sphere_Bfield(
     in the inside (see e.g. "Theoretical Physics, Bertelmann").
     """
     xp = array_namespace(observers, diameters, polarizations)
-    observers, diameters, polarizations = xp_promote(observers, diameters, polarizations, force_floating=True, xp=xp)
+    observers, diameters, polarizations = xp_promote(
+        observers, diameters, polarizations, force_floating=True, xp=xp
+    )
     r = xp.linalg.vector_norm(observers, axis=-1)
     r_sphere = xp.abs(diameters) / 2.0
 
@@ -90,8 +92,9 @@ def B_out(polarization, observers, r_sphere, r):
             / 3.0
         ).T
 
-    B = B_out(polarization, observers, r_sphere[xp.newaxis, :], r[xp.newaxis, :])
-    B = xpx.at(B)[mask_in].set(B_in(polarization[mask_in]))
+    B_o = B_out(polarization, observers, r_sphere[xp.newaxis, :], r[xp.newaxis, :])
+    B_i = B_in(polarization)
+    B = xp.where(mask_in, B_i, B_o)
 
     return B
 
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -1,8 +1,12 @@
 # here all core functions should be tested properly - ideally against FEM
 from __future__ import annotations
 
-import array_api_strict as xp
+import array_api_strict
+import dask.array
+import jax
 import numpy as np
+import pytest
+from array_api_extra.testing import lazy_xp_function, patch_lazy_xp_functions
 
 from magpylib.core import (
     current_circle_Hfield,
@@ -16,30 +20,46 @@
     triangle_Bfield,
 )
 
+lazy_xp_function(current_circle_Hfield)
+lazy_xp_function(current_polyline_Hfield)
+lazy_xp_function(dipole_Hfield)
+lazy_xp_function(magnet_cuboid_Bfield)
+lazy_xp_function(magnet_cylinder_axial_Bfield)
+lazy_xp_function(magnet_cylinder_diametral_Hfield)
+lazy_xp_function(magnet_cylinder_segment_Hfield)
+lazy_xp_function(magnet_sphere_Bfield)
+lazy_xp_function(triangle_Bfield)
 
-def test_magnet_sphere_Bfield():
+
+@pytest.fixture(params=[np, array_api_strict, jax.numpy, dask.array])
+def xp(request, monkeypatch):
+    patch_lazy_xp_functions(request, monkeypatch, xp=request.param)
+    return request.param
+
+
+def test_magnet_sphere_Bfield(xp):
     "magnet_sphere_Bfield test"
     B = magnet_sphere_Bfield(
         observers=xp.asarray([(0, 0, 0)]),
         diameters=xp.asarray([1]),
         polarizations=xp.asarray([(0, 0, 1)]),
     )
     Btest = xp.asarray([(0, 0, 2 / 3)])
-    np.testing.assert_allclose(B, Btest)
+    np.testing.assert_allclose(B, Btest, rtol=1e-4)
 
 
-def test_current_circle_Hfield():
+def test_current_circle_Hfield(xp):
     Htest = xp.asarray([[0.09098208, 0.09415448], [0.0, 0.0], [0.07677892, 0.22625335]])
     H = current_circle_Hfield(
         r0=xp.asarray([1, 2]),
         r=xp.asarray([1, 1]),
         z=xp.asarray([1, 2]),
         i0=xp.asarray([1, 3]),
     )
-    np.testing.assert_allclose(H, Htest)
+    np.testing.assert_allclose(H, Htest, rtol=1e-4)
 
 
-def test_current_polyline_Hfield():
+def test_current_polyline_Hfield(xp):
     Htest = xp.asarray([[0.0, -2.29720373, 2.29720373], [0.0, 0.59785204, -0.59785204]])
 
     H = current_polyline_Hfield(
@@ -48,10 +68,10 @@ def test_current_polyline_Hfield():
         segments_end=xp.asarray([(1, 0, 0), (-1, 0, 0)]),
         currents=xp.asarray([100, 200]),
     )
-    np.testing.assert_allclose(H, Htest)
+    np.testing.assert_allclose(H, Htest, rtol=1e-4)
 
 
-def test_dipole_Hfield():
+def test_dipole_Hfield(xp):
     Htest = xp.asarray(
         [
             [2.89501155e-13, 1.53146915e03, 1.53146915e03],
@@ -62,35 +82,35 @@ def test_dipole_Hfield():
         observers=xp.asarray([(1, 1, 1), (2, 2, 2)]),
         moments=xp.asarray([(1e5, 0, 0), (0, 0, 1e5)]),
     )
-    np.testing.assert_allclose(H, Htest)
+    np.testing.assert_allclose(H, Htest, atol=1e-6)
 
 
-def test_magnet_cuboid_Bfield():
+def test_magnet_cuboid_Bfield(xp):
     Btest = xp.asarray(
         [
-            [1.56103722e-02, 1.56103722e-02, -3.53394965e-17],
-            [7.73243250e-03, 6.54431406e-03, 1.04789520e-02],
+            [1.5610372220113404e-02, 1.5610372220113404e-02, -3.5339496460705743e-17],
+            [7.7324325000007145e-03, 6.5443140645650129e-03, 1.0478952028501879e-02],
         ]
     )
     B = magnet_cuboid_Bfield(
         observers=xp.asarray([(1, 1, 1), (2, 2, 2)]),
         dimensions=xp.asarray([(1, 1, 1), (1, 2, 3)]),
         polarizations=xp.asarray([(0, 0, 1), (0.5, 0.5, 0)]),
     )
-    np.testing.assert_allclose(B, Btest)
+    np.testing.assert_allclose(B, Btest, atol=1e-6)
 
 
-def test_magnet_cylinder_axial_Bfield():
+def test_magnet_cylinder_axial_Bfield(xp):
     Btest = xp.asarray([[0.05561469, 0.04117919], [0.0, 0.0], [0.06690167, 0.01805674]])
     B = magnet_cylinder_axial_Bfield(
         z0=xp.asarray([1, 2]),
         r=xp.asarray([1, 2]),
         z=xp.asarray([2, 3]),
     )
-    np.testing.assert_allclose(B, Btest)
+    np.testing.assert_allclose(B, Btest, rtol=1e-4)
 
 
-def test_magnet_cylinder_diametral_Hfield():
+def test_magnet_cylinder_diametral_Hfield(xp):
     Btest = xp.asarray(
         [
             [-0.020742122169014, 0.007307203574376],
@@ -104,10 +124,10 @@ def test_magnet_cylinder_diametral_Hfield():
         z=xp.asarray([2, 3]),
         phi=xp.asarray([0.1, np.pi / 4]),
     )
-    np.testing.assert_allclose(B, Btest)
+    np.testing.assert_allclose(B, Btest, rtol=1e-4)
 
 
-def test_magnet_cylinder_segment_Hfield():
+def test_magnet_cylinder_segment_Hfield(xp):
     Btest = xp.asarray(
         [
             [-1948.14367497, 32319.94437208, 17616.88571231],
@@ -119,10 +139,10 @@ def test_magnet_cylinder_segment_Hfield():
         dimensions=xp.asarray([(1, 2, 0.1, 0.2, -1, 1), (1, 2, 0.3, 0.9, 0, 1)]),
         magnetizations=xp.asarray([(1e7, 0.1, 0.2), (1e6, 1.1, 2.2)]),
     )
-    np.testing.assert_allclose(B, Btest)
+    np.testing.assert_allclose(B, Btest, rtol=1e-4)
 
 
-def test_triangle_Bfield():
+def test_triangle_Bfield(xp):
     Btest = xp.asarray(
         [[7.45158965, 4.61994866, 3.13614132], [2.21345618, 2.67710148, 2.21345618]]
     )
@@ -136,4 +156,4 @@ def test_triangle_Bfield():
         ),
         polarizations=xp.asarray([(1.0, 1.0, 1.0), (1.0, 1.0, 0.0)]) * 1e3,
     )
-    np.testing.assert_allclose(B, Btest)
+    np.testing.assert_allclose(B, Btest, rtol=1e-4)
