Adds array api support to cylinder segment BHMJ level

purepani · purepani · commit bdb83ca2cd65 · 2025-05-20T16:32:31.000-05:00
diff --git a/src/magpylib/_src/fields/field_BH_cylinder_segment.py b/src/magpylib/_src/fields/field_BH_cylinder_segment.py
@@ -2375,42 +2375,47 @@ def BHJM_cylinder_segment(
     - translate cylinder segment field to BHJM
     - special cases catching
     """
+    xp = array_namespace(observers, dimension, polarization)
     check_field_input(field)
 
-    BHJM = polarization.astype(float)
+    observers = xp.astype(observers, xp.float64)
+    dimension = xp.astype(dimension, xp.float64)
+    polarization = xp.astype(polarization, xp.float64)
 
-    r1, r2, h, phi1, phi2 = dimension.T
-    r1 = abs(r1)
-    r2 = abs(r2)
-    h = abs(h)
+    BHJM = xp.astype(polarization, xp.float64)
+
+    # r1, r2, h, phi1, phi2 = dimension.T
+    r1 = xp.abs(dimension[:, 0])
+    r2 = xp.abs(dimension[:, 1])
+    h = xp.abs(dimension[:, 2])
     z1, z2 = -h / 2, h / 2
 
     # transform dim deg->rad
-    phi1 = phi1 / 180 * np.pi
-    phi2 = phi2 / 180 * np.pi
-    dim = np.array([r1, r2, phi1, phi2, z1, z2]).T
+    phi1 = dimension[:, 3] / 180 * xp.pi
+    phi2 = dimension[:, 4] / 180 * xp.pi
+    dim = xp.asarray([r1, r2, phi1, phi2, z1, z2]).T
 
     # transform obs_pos to Cy CS --------------------------------------------
-    x, y, z = observers.T
-    r, phi = np.sqrt(x**2 + y**2), np.arctan2(y, x)
-    pos_obs_cy = np.concatenate(((r,), (phi,), (z,)), axis=0).T
+    x, y, z = (observers[:, i] for i in range(3))
+    r, phi = xp.sqrt(x**2 + y**2), xp.atan2(y, x)
+    pos_obs_cy = xp.concat((r[None, ...], phi[None, ...], z[None, ...]), axis=0).T
 
     # determine when points lie inside and on surface of magnet -------------
 
     # mask_inside = None
     # if in_out == "auto":
     # phip1 in [-2pi,0], phio2 in [0,2pi]
     phio1 = phi
-    phio2 = phi - np.sign(phi) * 2 * np.pi
+    phio2 = phi - xp.sign(phi) * 2 * xp.pi
 
     # phi=phi1, phi=phi2
     mask_phi1 = close(phio1, phi1) | close(phio2, phi1)
     mask_phi2 = close(phio1, phi2) | close(phio2, phi2)
 
     # r, phi ,z lies in-between, avoid numerical fluctuations (e.g. due to rotations) by including 1e-14
     mask_r_in = (r1 - 1e-14 < r) & (r < r2 + 1e-14)
-    mask_phi_in = (np.sign(phio1 - phi1) != np.sign(phio1 - phi2)) | (
-        np.sign(phio2 - phi1) != np.sign(phio2 - phi2)
+    mask_phi_in = (xp.sign(phio1 - phi1) != xp.sign(phio1 - phi2)) | (
+        xp.sign(phio2 - phi1) != xp.sign(phio2 - phi2)
     )
     mask_z_in = (z1 - 1e-14 < z) & (z < z2 + 1e-14)
 
@@ -2425,13 +2430,13 @@ def BHJM_cylinder_segment(
     # inside
     mask_inside = mask_r_in & mask_phi_in & mask_z_in
     # else:
-    #     mask_inside = np.full(len(observers), in_out == "inside")
-    #     mask_not_on_surf = np.full(len(observers), True)
+    #     mask_inside = xp.full((observers.shape[0]), in_out == "inside")
+    #     mask_not_on_surf = xp.full((observers.shape[0]), True)
     # WARNING @alex
     #   1. inside and not_on_surface are not the same! Can't just put to true.
 
     # return 0 when all points are on surface
-    if not np.any(mask_not_on_surf):
+    if not xp.any(mask_not_on_surf):
         return BHJM * 0
 
     if field == "J":
@@ -2451,19 +2456,21 @@ def BHJM_cylinder_segment(
     phi = phi[mask_not_on_surf]
 
     # transform mag to spherical CS -----------------------------------------
-    m = np.sqrt(pol[:, 0] ** 2 + pol[:, 1] ** 2 + pol[:, 2] ** 2) / MU0  # J -> M
-    phi_m = np.arctan2(pol[:, 1], pol[:, 0])
-    th_m = np.arctan2(np.sqrt(pol[:, 0] ** 2 + pol[:, 1] ** 2), pol[:, 2])
-    mag_sph = np.concatenate(((m,), (phi_m,), (th_m,)), axis=0).T
+    m = xp.sqrt(pol[:, 0] ** 2 + pol[:, 1] ** 2 + pol[:, 2] ** 2) / MU0  # J -> M
+    phi_m = xp.atan2(pol[:, 1], pol[:, 0])
+    th_m = xp.atan2(xp.sqrt(pol[:, 0] ** 2 + pol[:, 1] ** 2), pol[:, 2])
+    mag_sph = xp.concat((m[None, ...], phi_m[None, ...], th_m[None, ...]), axis=0).T
 
     # compute H and transform to cart CS -------------------------------------
     H_cy = magnet_cylinder_segment_Hfield(
         magnetizations=mag_sph, dimensions=dim, observers=pos_obs_cy
     )
-    Hr, Hphi, Hz = H_cy.T
-    Hx = Hr * np.cos(phi) - Hphi * np.sin(phi)
-    Hy = Hr * np.sin(phi) + Hphi * np.cos(phi)
-    BHJM[mask_not_on_surf] = np.concatenate(((Hx,), (Hy,), (Hz,)), axis=0).T
+    Hr, Hphi, Hz = H_cy[:, 0], H_cy[:, 1], H_cy[:, 2]
+    Hx = Hr * xp.cos(phi) - Hphi * xp.sin(phi)
+    Hy = Hr * xp.sin(phi) + Hphi * xp.cos(phi)
+    BHJM[mask_not_on_surf] = xp.concat(
+        (Hx[None, ...], Hy[None, ...], Hz[None, ...]), axis=0
+    ).T
 
     if field == "H":
         return BHJM
@@ -2484,3 +2491,25 @@ def BHJM_cylinder_segment(
     #     field_values=H_all,
     #     mask_inside=mask_inside & mask_not_on_surf,
     # )
+
+    # return convert_HBMJ(
+    #     output_field_type=field,
+    #     polarization=polarization,
+    #     input_field_type="H",
+    #     field_values=H_all,
+    #     mask_inside=mask_inside & mask_not_on_surf,
+    # )
+
+    # return convert_HBMJ(
+    #     output_field_type=field,
+    #     polarization=polarization,
+    #     input_field_type="H",
+    #     field_values=H_all,
+    #     mask_inside=mask_inside & mask_not_on_surf,
+    # )
+    #     output_field_type=field,
+    #     polarization=polarization,
+    #     input_field_type="H",
+    #     field_values=H_all,
+    #     mask_inside=mask_inside & mask_not_on_surf,
+    # )
diff --git a/tests/test_BHMJ_level.py b/tests/test_BHMJ_level.py
@@ -214,23 +214,23 @@ def test_BHJM_magnet_sphere():
 
 def test_field_cylinder_segment_BH():
     """CylinderSegment field test"""
-    pol = np.array(
+    pol = xp.asarray(
         [
             (0, 0, 0),
             (1, 2, 3),
             (2, 3, -1),
             (2, 3, -1),
         ]
     )
-    dim = np.array(
+    dim = xp.asarray(
         [
             (1, 2, 3, 10, 20),
             (1, 2, 3, 10, 20),
             (1, 3, 2, -50, 50),
             (0.1, 5, 2, 20, 370),
         ]
     )
-    obs = np.array(
+    obs = xp.asarray(
         [
             (1, 2, 3),
             (1, -1, 0),

Original file line number	Diff line number	Diff line change
`@@ -214,23 +214,23 @@ def test_BHJM_magnet_sphere():`
`214`	`214`
`215`	`215`	`def test_field_cylinder_segment_BH():`
`216`	`216`	`"""CylinderSegment field test"""`
`217`		`- pol = np.array(`
	`217`	`+ pol = xp.asarray(`
`218`	`218`	`[`
`219`	`219`	`(0, 0, 0),`
`220`	`220`	`(1, 2, 3),`
`221`	`221`	`(2, 3, -1),`
`222`	`222`	`(2, 3, -1),`
`223`	`223`	`]`
`224`	`224`	`)`
`225`		`- dim = np.array(`
	`225`	`+ dim = xp.asarray(`
`226`	`226`	`[`
`227`	`227`	`(1, 2, 3, 10, 20),`
`228`	`228`	`(1, 2, 3, 10, 20),`
`229`	`229`	`(1, 3, 2, -50, 50),`
`230`	`230`	`(0.1, 5, 2, 20, 370),`
`231`	`231`	`]`
`232`	`232`	`)`
`233`		`- obs = np.array(`
	`233`	`+ obs = xp.asarray(`
`234`	`234`	`[`
`235`	`235`	`(1, 2, 3),`
`236`	`236`	`(1, -1, 0),`