diff --git a/pyproject.toml b/pyproject.toml index 36301bf4a..e001243fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,9 @@ dependencies = [ "scipy>=1.8", "matplotlib>=3.6", "plotly>=5.16", + "array-api-compat>=1.11.2", + "array-api-extra>=0.7.1", + "typing-extensions>=4.13.2", ] [dependency-groups] @@ -68,6 +71,10 @@ test = [ "imageio[tifffile,ffmpeg]", "jupyterlab", "anywidget", + "array-api-strict>=2.3.1", + "jax", + "numpy", + "dask[array]", ] binder = [ "jupytext", diff --git a/src/magpylib/_src/array_api_utils.py b/src/magpylib/_src/array_api_utils.py new file mode 100644 index 000000000..2811ef457 --- /dev/null +++ b/src/magpylib/_src/array_api_utils.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +from types import ModuleType +from typing import Any, Literal, TypeAlias + +import numpy as np +import numpy.typing as npt +from array_api_compat import ( + array_namespace, +) +from array_api_compat import ( + is_numpy_namespace as is_numpy, +) +from array_api_compat import ( + is_torch_namespace as is_torch, +) + +Array: TypeAlias = Any # To be changed to a Protocol later (see array-api#589) +ArrayLike: TypeAlias = Array | npt.ArrayLike + + +def _check_finite(array: Array, xp: ModuleType) -> None: + """Check for NaNs or Infs.""" + if not xp.all(xp.isfinite(array)): + msg = "array must not contain infs or NaNs" + raise ValueError(msg) + + +def _asarray( + array: ArrayLike, + dtype: Any = None, + order: Literal["K", "A", "C", "F"] | None = None, + copy: bool | None = None, + *, + xp: ModuleType | None = None, + check_finite: bool = False, + subok: bool = False, +) -> Array: + """SciPy-specific replacement for `np.asarray` with `order`, `check_finite`, and + `subok`. + + Memory layout parameter `order` is not exposed in the Array API standard. + `order` is only enforced if the input array implementation + is NumPy based, otherwise `order` is just silently ignored. + + `check_finite` is also not a keyword in the array API standard; included + here for convenience rather than that having to be a separate function + call inside SciPy functions. + + `subok` is included to allow this function to preserve the behaviour of + `np.asanyarray` for NumPy based inputs. + """ + if xp is None: + xp = array_namespace(array) + if is_numpy(xp): + # Use NumPy API to support order + if copy is True: + array = np.array(array, order=order, dtype=dtype, subok=subok) + elif subok: + array = np.asanyarray(array, order=order, dtype=dtype) + else: + array = np.asarray(array, order=order, dtype=dtype) + else: + try: + array = xp.asarray(array, dtype=dtype, copy=copy) + except TypeError: + coerced_xp = array_namespace(xp.asarray(3)) + array = coerced_xp.asarray(array, dtype=dtype, copy=copy) + + if check_finite: + _check_finite(array, xp) + + return array + + +def xp_default_dtype(xp): + """Query the namespace-dependent default floating-point dtype.""" + if is_torch(xp): + # historically, we allow pytorch to keep its default of float32 + return xp.get_default_dtype() + # we default to float64 + return xp.float64 + + +def xp_result_type(*args, force_floating=False, xp): + """ + Returns the dtype that results from applying type promotion rules + (see Array API Standard Type Promotion Rules) to the arguments. Augments + standard `result_type` in a few ways: + + - There is a `force_floating` argument that ensures that the result type + is floating point, even when all args are integer. + - When a TypeError is raised (e.g. due to an unsupported promotion) + and `force_floating=True`, we define a custom rule: use the result type + of the default float and any other floats passed. See + https://github.com/scipy/scipy/pull/22695/files#r1997905891 + for rationale. + - This function accepts array-like iterables, which are immediately converted + to the namespace's arrays before result type calculation. Consequently, the + result dtype may be different when an argument is `1.` vs `[1.]`. + + Typically, this function will be called shortly after `array_namespace` + on a subset of the arguments passed to `array_namespace`. + """ + args = [ + (_asarray(arg, subok=True, xp=xp) if np.iterable(arg) else arg) for arg in args + ] + args_not_none = [arg for arg in args if arg is not None] + if force_floating: + args_not_none.append(1.0) + + if is_numpy(xp) and xp.__version__ < "2.0": + # Follow NEP 50 promotion rules anyway + args_not_none = [ + arg.dtype if getattr(arg, "size", 0) == 1 else arg for arg in args_not_none + ] + return xp.result_type(*args_not_none) + + try: # follow library's preferred promotion rules + return xp.result_type(*args_not_none) + except TypeError: # mixed type promotion isn't defined + if not force_floating: + raise + # use `result_type` of default floating point type and any floats present + # This can be revisited, but right now, the only backends that get here + # are array-api-strict (which is not for production use) and PyTorch + # (due to data-apis/array-api-compat#279). + float_args = [] + for arg in args_not_none: + arg_array = xp.asarray(arg) if np.isscalar(arg) else arg + dtype = getattr(arg_array, "dtype", arg) + if xp.isdtype(dtype, ("real floating", "complex floating")): + float_args.append(arg) + return xp.result_type(*float_args, xp_default_dtype(xp)) + + +def xp_promote(*args, broadcast=False, force_floating=False, xp): + """ + Promotes elements of *args to result dtype, ignoring `None`s. + Includes options for forcing promotion to floating point and + broadcasting the arrays, again ignoring `None`s. + Type promotion rules follow `xp_result_type` instead of `xp.result_type`. + + Typically, this function will be called shortly after `array_namespace` + on a subset of the arguments passed to `array_namespace`. + + This function accepts array-like iterables, which are immediately converted + to the namespace's arrays before result type calculation. Consequently, the + result dtype may be different when an argument is `1.` vs `[1.]`. + + See Also + -------- + xp_result_type + """ + args = [ + (_asarray(arg, subok=True, xp=xp) if np.iterable(arg) else arg) for arg in args + ] # solely to prevent double conversion of iterable to array + + dtype = xp_result_type(*args, force_floating=force_floating, xp=xp) + + args = [ + (_asarray(arg, dtype=dtype, subok=True, xp=xp) if arg is not None else arg) + for arg in args + ] + + if not broadcast: + return args[0] if len(args) == 1 else tuple(args) + + args_not_none = [arg for arg in args if arg is not None] + + # determine result shape + shapes = {arg.shape for arg in args_not_none} + try: + shape = ( + np.broadcast_shapes(*shapes) if len(shapes) != 1 else args_not_none[0].shape + ) + except ValueError as e: + message = "Array shapes are incompatible for broadcasting." + raise ValueError(message) from e + + out = [] + for arg in args: + if arg is None: + out.append(arg) + continue + + # broadcast only if needed + # Even if two arguments need broadcasting, this is faster than + # `broadcast_arrays`, especially since we've already determined `shape` + if arg.shape != shape: + kwargs = {"subok": True} if is_numpy(xp) else {} + arg = xp.broadcast_to(arg, shape, **kwargs) + + # This is much faster than xp.astype(arg, dtype, copy=False) + if arg.dtype != dtype: + arg = xp.astype(arg, dtype) + + out.append(arg) + + return out[0] if len(out) == 1 else tuple(out) diff --git a/src/magpylib/_src/fields/field_BH_polyline.py b/src/magpylib/_src/fields/field_BH_polyline.py index 7d998a321..06b98b630 100644 --- a/src/magpylib/_src/fields/field_BH_polyline.py +++ b/src/magpylib/_src/fields/field_BH_polyline.py @@ -5,10 +5,12 @@ # pylint: disable=too-many-positional-arguments from __future__ import annotations +import array_api_extra as xpx import numpy as np -from numpy.linalg import norm +from array_api_compat import array_namespace from scipy.constants import mu_0 as MU0 +from magpylib._src.array_api_utils import cross, xp_promote from magpylib._src.input_checks import check_field_input @@ -124,72 +126,99 @@ def current_polyline_Hfield( Be careful with magnetic fields of discontinued segments. They are unphysical and can lead to unphysical effects. """ + xp = array_namespace(observers, segments_start, segments_end, currents) + observers, segments_start, segments_end, currents = xp_promote( + observers, segments_start, segments_end, currents, force_floating=True, xp=xp + ) + dtype = observers.dtype # rename p1, p2, po = segments_start, segments_end, observers - # make dimensionless (avoid all large/small input problems) by introducing # the segment length as characteristic length scale. - norm_12 = norm(p1 - p2, axis=1) + norm_12 = xp.linalg.vector_norm(p1 - p2, axis=-1) p1 = (p1.T / norm_12).T p2 = (p2.T / norm_12).T po = (po.T / norm_12).T # p4 = projection of pos_obs onto line p1-p2 - t = np.sum((po - p1) * (p1 - p2), axis=1) + t = xp.sum((po - p1) * (p1 - p2), axis=1) p4 = p1 + (t * (p1 - p2).T).T # distance of observers from line - norm_o4 = norm(po - p4, axis=1) + norm_o4 = xp.linalg.vector_norm(po - p4, axis=-1) # separate on-line cases (-> B=0) mask1 = norm_o4 < 1e-15 # account for numerical issues # continue only with general off-line cases - if np.any(mask1): - not_mask1 = ~mask1 - po = po[not_mask1] - p1 = p1[not_mask1] - p2 = p2[not_mask1] - p4 = p4[not_mask1] - norm_12 = norm_12[not_mask1] - norm_o4 = norm_o4[not_mask1] - currents = currents[not_mask1] - - # determine field direction - cros_ = np.cross(p2 - p1, po - p4) - norm_cros = norm(cros_, axis=1) - eB = (cros_.T / norm_cros).T - - # compute angles - norm_o1 = norm( - po - p1, axis=1 - ) # improve performance by computing all norms at once - norm_o2 = norm(po - p2, axis=1) - norm_41 = norm(p4 - p1, axis=1) - norm_42 = norm(p4 - p2, axis=1) - sinTh1 = norm_41 / norm_o1 - sinTh2 = norm_42 / norm_o2 - deltaSin = np.empty((len(po),)) - - # determine how p1,p2,p4 are sorted on the line (to get sinTH signs) - # both points below - mask2 = (norm_41 > 1) * (norm_41 > norm_42) - deltaSin[mask2] = abs(sinTh1[mask2] - sinTh2[mask2]) - # both points above - mask3 = (norm_42 > 1) * (norm_42 > norm_41) - deltaSin[mask3] = abs(sinTh2[mask3] - sinTh1[mask3]) - # one above one below or one equals p4 - mask4 = ~mask2 * ~mask3 - deltaSin[mask4] = abs(sinTh1[mask4] + sinTh2[mask4]) - - # B = (deltaSin / norm_o4 * eB.T / norm_12 * current * 1e-7).T - - # avoid array creation if possible - if np.any(mask1): - H = np.zeros_like(observers, dtype=float) - H[~mask1] = (deltaSin / norm_o4 * eB.T / norm_12 * currents / (4 * np.pi)).T - return H - return (deltaSin / norm_o4 * eB.T / norm_12 * currents / (4 * np.pi)).T + def H_calc(po, p1, p2, p4, norm_12, norm_o4, currents): + cros_ = cross(p2 - p1, po - p4) + + norm_cros = xp.linalg.vector_norm(cros_, axis=-1) + eB = (cros_.T / norm_cros).T + + # compute angles + norm_o1 = xp.linalg.vector_norm( + po - p1, axis=-1 + ) # improve performance by computing all norms at once + norm_o2 = xp.linalg.vector_norm(po - p2, axis=-1) + norm_41 = xp.linalg.vector_norm(p4 - p1, axis=-1) + norm_42 = xp.linalg.vector_norm(p4 - p2, axis=-1) + sinTh1 = norm_41 / norm_o1 + sinTh2 = norm_42 / norm_o2 + + mask2 = (norm_41 > 1) & (norm_41 > norm_42) + mask3 = (norm_42 > 1) & (norm_42 > norm_41) + mask4 = ~mask2 & ~mask3 + + deltaSin = xpx.apply_where( + mask2, + (sinTh1, sinTh2), + lambda sinTh1, sinTh2: sinTh1 - sinTh2, + fill_value=0.0, + ) + deltaSin = xpx.apply_where( + mask3, + (sinTh1, sinTh2, deltaSin), + lambda sinTh1, sinTh2, deltaSin: sinTh2 - sinTh1, + lambda sinTh1, sinTh2, deltaSin: deltaSin, + ) + deltaSin = xpx.apply_where( + mask4, + (sinTh1, sinTh2, deltaSin), + lambda sinTh1, sinTh2, deltaSin: sinTh1 + sinTh2, + lambda sinTh1, sinTh2, deltaSin: deltaSin, + ) + + # deltaSin = xp.empty((po.shape[0],)) + + # determine how p1,p2,p4 are sorted on the line (to get sinTH signs) + # both points below + # mask2 = (norm_41 > 1) & (norm_41 > norm_42) + # deltaSin = xpx.at(deltaSin)[mask2].set(xp.abs(sinTh1[mask2] - sinTh2[mask2])) + # both points above + # mask3 = (norm_42 > 1) & (norm_42 > norm_41) + # deltaSin = xpx.at(deltaSin)[mask3].set(xp.abs(sinTh2[mask3] - sinTh1[mask3])) + # one above one below or one equals p4 + # mask4 = ~mask2 & ~mask3 + # deltaSin = xpx.at(deltaSin)[mask4].set(xp.abs(sinTh1[mask4] + sinTh2[mask4])) + + # B = (deltaSin / norm_o4 * eB.T / norm_12 * current * 1e-7).T + + return (deltaSin / norm_o4 * eB.T / norm_12 * currents / (4 * np.pi)).T + + # H = xpx.apply_where( + # ~mask1[:, None], + # (po, p1, p2, p4, norm_12[:, None], norm_o4[:, None], currents[:, None]), + # H_calc, + # fill_value=0.0, + # ) + H = xp.where( + ~mask1[:, None], + H_calc(po, p1, p2, p4, norm_12, norm_o4, currents), + 0.0, + ) + return H def BHJM_current_polyline( diff --git a/tests/test_core.py b/tests/test_core.py index 9879cb142..d9408d7b8 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1,17 +1,159 @@ # here all core functions should be tested properly - ideally against FEM from __future__ import annotations +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._src.fields.field_BH_sphere import magnet_sphere_Bfield +from magpylib.core import ( + current_circle_Hfield, + current_polyline_Hfield, + dipole_Hfield, + magnet_cuboid_Bfield, + magnet_cylinder_axial_Bfield, + magnet_cylinder_diametral_Hfield, + magnet_cylinder_segment_Hfield, + magnet_sphere_Bfield, + 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=np.array([(0, 0, 0)]), - diameters=np.array([1]), - polarizations=np.array([(0, 0, 1)]), + 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, rtol=1e-4) + + +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, rtol=1e-4) + + +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( + observers=xp.asarray([(1, 1, 1), (2, 2, 2)]), + segments_start=xp.asarray([(0, 0, 0), (0, 0, 0)]), + segments_end=xp.asarray([(1, 0, 0), (-1, 0, 0)]), + currents=xp.asarray([100, 200]), + ) + np.testing.assert_allclose(H, Htest, rtol=1e-4) + + +def test_dipole_Hfield(xp): + Htest = xp.asarray( + [ + [2.89501155e-13, 1.53146915e03, 1.53146915e03], + [1.91433644e02, 1.91433644e02, 3.61876444e-14], + ] + ) + H = 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, atol=1e-6) + + +def test_magnet_cuboid_Bfield(xp): + Btest = xp.asarray( + [ + [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, atol=1e-6) + + +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, rtol=1e-4) + + +def test_magnet_cylinder_diametral_Hfield(xp): + Btest = xp.asarray( + [ + [-0.020742122169014, 0.007307203574376], + [0.004597868528024, 0.020075245863212], + [0.05533684822464, 0.029118084290573], + ], + ) + B = magnet_cylinder_diametral_Hfield( + z0=xp.asarray([1, 2]), + r=xp.asarray([1, 2]), + z=xp.asarray([2, 3]), + phi=xp.asarray([0.1, np.pi / 4]), + ) + np.testing.assert_allclose(B, Btest, rtol=1e-4) + + +def test_magnet_cylinder_segment_Hfield(xp): + Btest = xp.asarray( + [ + [-1948.14367497, 32319.94437208, 17616.88571231], + [14167.64961763, 1419.94126065, 17921.6463117], + ] + ) + B = magnet_cylinder_segment_Hfield( + observers=xp.asarray([(1, 1, 2), (0, 0, 0)]), + 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, rtol=1e-4) + + +def test_triangle_Bfield(xp): + Btest = xp.asarray( + [[7.45158965, 4.61994866, 3.13614132], [2.21345618, 2.67710148, 2.21345618]] + ) + B = triangle_Bfield( + observers=xp.asarray([(2.0, 1.0, 1.0), (2.0, 2.0, 2.0)]), + vertices=xp.asarray( + [ + [(0.0, 0.0, 0.0), (0.0, 0.0, 1.0), (1.0, 0.0, 0.0)], + [(0.0, 0.0, 0.0), (0.0, 0.0, 1.0), (1.0, 0.0, 0.0)], + ] + ), + polarizations=xp.asarray([(1.0, 1.0, 1.0), (1.0, 1.0, 0.0)]) * 1e3, ) - Btest = np.array([(0, 0, 2 / 3)]) - np.testing.assert_allclose(B, Btest) + np.testing.assert_allclose(B, Btest, rtol=1e-4)