From 27140f221a70ced1ee04fd23fce37466a0e5b3b8 Mon Sep 17 00:00:00 2001 From: purepani Date: Tue, 20 May 2025 15:43:24 -0500 Subject: [PATCH 01/16] Adds core tests --- tests/test_core.py | 120 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 119 insertions(+), 1 deletion(-) diff --git a/tests/test_core.py b/tests/test_core.py index 9879cb142..d72781842 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -3,7 +3,17 @@ import numpy as np -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, +) def test_magnet_sphere_Bfield(): @@ -15,3 +25,111 @@ def test_magnet_sphere_Bfield(): ) Btest = np.array([(0, 0, 2 / 3)]) np.testing.assert_allclose(B, Btest) + + +def test_current_circle_Hfield(): + Htest = np.array([[0.09098208, 0.09415448], [0.0, 0.0], [0.07677892, 0.22625335]]) + H = current_circle_Hfield( + r0=np.array([1, 2]), + r=np.array([1, 1]), + z=np.array([1, 2]), + i0=np.array([1, 3]), + ) + np.testing.assert_allclose(H, Htest) + + +def test_current_polyline_Hfield(): + Htest = np.array([[0.0, -2.29720373, 2.29720373], [0.0, 0.59785204, -0.59785204]]) + + H = current_polyline_Hfield( + observers=np.array([(1, 1, 1), (2, 2, 2)]), + segments_start=np.array([(0, 0, 0), (0, 0, 0)]), + segments_end=np.array([(1, 0, 0), (-1, 0, 0)]), + currents=np.array([100, 200]), + ) + np.testing.assert_allclose(H, Htest) + + +def test_dipole_Hfield(): + Htest = np.array( + [ + [2.89501155e-13, 1.53146915e03, 1.53146915e03], + [1.91433644e02, 1.91433644e02, 3.61876444e-14], + ] + ) + H = dipole_Hfield( + observers=np.array([(1, 1, 1), (2, 2, 2)]), + moments=np.array([(1e5, 0, 0), (0, 0, 1e5)]), + ) + np.testing.assert_allclose(H, Htest) + + +def test_magnet_cuboid_Bfield(): + Btest = np.array( + [ + [1.56103722e-02, 1.56103722e-02, -3.53394965e-17], + [7.73243250e-03, 6.54431406e-03, 1.04789520e-02], + ] + ) + B = magnet_cuboid_Bfield( + observers=np.array([(1, 1, 1), (2, 2, 2)]), + dimensions=np.array([(1, 1, 1), (1, 2, 3)]), + polarizations=np.array([(0, 0, 1), (0.5, 0.5, 0)]), + ) + np.testing.assert_allclose(B, Btest) + + +def test_magnet_cylinder_axial_Bfield(): + Btest = np.array([[0.05561469, 0.04117919], [0.0, 0.0], [0.06690167, 0.01805674]]) + B = magnet_cylinder_axial_Bfield( + z0=np.array([1, 2]), + r=np.array([1, 2]), + z=np.array([2, 3]), + ) + np.testing.assert_allclose(B, Btest) + + +def test_magnet_cylinder_diametral_Hfield(): + Btest = np.array( + [ + [-0.020742122169014, 0.007307203574376], + [0.004597868528024, 0.020075245863212], + [0.05533684822464, 0.029118084290573], + ], + ) + B = magnet_cylinder_diametral_Hfield( + z0=np.array([1, 2]), + r=np.array([1, 2]), + z=np.array([2, 3]), + phi=np.array([0.1, np.pi / 4]), + ) + np.testing.assert_allclose(B, Btest) + + +def test_magnet_cylinder_segment_Hfield(): + Btest = np.array( + [ + [-1948.14367497, 32319.94437208, 17616.88571231], + [14167.64961763, 1419.94126065, 17921.6463117], + ] + ) + B = magnet_cylinder_segment_Hfield( + observers=np.array([(1, 1, 2), (0, 0, 0)]), + dimensions=np.array([(1, 2, 0.1, 0.2, -1, 1), (1, 2, 0.3, 0.9, 0, 1)]), + magnetizations=np.array([(1e7, 0.1, 0.2), (1e6, 1.1, 2.2)]), + ) + np.testing.assert_allclose(B, Btest) + + +def test_triangle_Bfield(): + Btest = np.array( + [[7.45158965, 4.61994866, 3.13614132], [2.21345618, 2.67710148, 2.21345618]] + ) + B = triangle_Bfield( + observers=np.array([(2, 1, 1), (2, 2, 2)]), + vertices=np.array( + [[(0, 0, 0), (0, 0, 1), (1, 0, 0)], [(0, 0, 0), (0, 0, 1), (1, 0, 0)]] + ), + polarizations=np.array([(1, 1, 1), (1, 1, 0)]) * 1e3, + ) + np.testing.assert_allclose(B, Btest) From addd3346acaaa38da5ef01f7145b81966770157d Mon Sep 17 00:00:00 2001 From: purepani Date: Tue, 6 May 2025 08:16:30 -0500 Subject: [PATCH 02/16] dev From 5db3a42d8a42e1c348346653557edca7a714cb21 Mon Sep 17 00:00:00 2001 From: purepani Date: Mon, 19 May 2025 15:56:16 -0500 Subject: [PATCH 03/16] Adds array_api_strict dependency to test --- pyproject.toml | 3 ++ tests/test_core.py | 82 ++++++++++++++++++++++++---------------------- 2 files changed, 46 insertions(+), 39 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 36301bf4a..3b6f44814 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,8 @@ dependencies = [ "scipy>=1.8", "matplotlib>=3.6", "plotly>=5.16", + "array-api-compat>=1.11.2", + "array-api-extra>=0.7.1", ] [dependency-groups] @@ -68,6 +70,7 @@ test = [ "imageio[tifffile,ffmpeg]", "jupyterlab", "anywidget", + "array-api-strict>=2.3.1", ] binder = [ "jupytext", diff --git a/tests/test_core.py b/tests/test_core.py index d72781842..d4585f740 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1,6 +1,7 @@ # here all core functions should be tested properly - ideally against FEM from __future__ import annotations +import array_api_strict as xp import numpy as np from magpylib.core import ( @@ -19,78 +20,78 @@ def test_magnet_sphere_Bfield(): "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 = np.array([(0, 0, 2 / 3)]) + Btest = xp.asarray([(0, 0, 2 / 3)]) np.testing.assert_allclose(B, Btest) def test_current_circle_Hfield(): - Htest = np.array([[0.09098208, 0.09415448], [0.0, 0.0], [0.07677892, 0.22625335]]) + Htest = xp.asarray([[0.09098208, 0.09415448], [0.0, 0.0], [0.07677892, 0.22625335]]) H = current_circle_Hfield( - r0=np.array([1, 2]), - r=np.array([1, 1]), - z=np.array([1, 2]), - i0=np.array([1, 3]), + 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) def test_current_polyline_Hfield(): - Htest = np.array([[0.0, -2.29720373, 2.29720373], [0.0, 0.59785204, -0.59785204]]) + Htest = xp.asarray([[0.0, -2.29720373, 2.29720373], [0.0, 0.59785204, -0.59785204]]) H = current_polyline_Hfield( - observers=np.array([(1, 1, 1), (2, 2, 2)]), - segments_start=np.array([(0, 0, 0), (0, 0, 0)]), - segments_end=np.array([(1, 0, 0), (-1, 0, 0)]), - currents=np.array([100, 200]), + 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) def test_dipole_Hfield(): - Htest = np.array( + Htest = xp.asarray( [ [2.89501155e-13, 1.53146915e03, 1.53146915e03], [1.91433644e02, 1.91433644e02, 3.61876444e-14], ] ) H = dipole_Hfield( - observers=np.array([(1, 1, 1), (2, 2, 2)]), - moments=np.array([(1e5, 0, 0), (0, 0, 1e5)]), + observers=xp.asarray([(1, 1, 1), (2, 2, 2)]), + moments=xp.asarray([(1e5, 0, 0), (0, 0, 1e5)]), ) np.testing.assert_allclose(H, Htest) def test_magnet_cuboid_Bfield(): - Btest = np.array( + Btest = xp.asarray( [ [1.56103722e-02, 1.56103722e-02, -3.53394965e-17], [7.73243250e-03, 6.54431406e-03, 1.04789520e-02], ] ) B = magnet_cuboid_Bfield( - observers=np.array([(1, 1, 1), (2, 2, 2)]), - dimensions=np.array([(1, 1, 1), (1, 2, 3)]), - polarizations=np.array([(0, 0, 1), (0.5, 0.5, 0)]), + 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) def test_magnet_cylinder_axial_Bfield(): - Btest = np.array([[0.05561469, 0.04117919], [0.0, 0.0], [0.06690167, 0.01805674]]) + Btest = xp.asarray([[0.05561469, 0.04117919], [0.0, 0.0], [0.06690167, 0.01805674]]) B = magnet_cylinder_axial_Bfield( - z0=np.array([1, 2]), - r=np.array([1, 2]), - z=np.array([2, 3]), + z0=xp.asarray([1, 2]), + r=xp.asarray([1, 2]), + z=xp.asarray([2, 3]), ) np.testing.assert_allclose(B, Btest) def test_magnet_cylinder_diametral_Hfield(): - Btest = np.array( + Btest = xp.asarray( [ [-0.020742122169014, 0.007307203574376], [0.004597868528024, 0.020075245863212], @@ -98,38 +99,41 @@ def test_magnet_cylinder_diametral_Hfield(): ], ) B = magnet_cylinder_diametral_Hfield( - z0=np.array([1, 2]), - r=np.array([1, 2]), - z=np.array([2, 3]), - phi=np.array([0.1, np.pi / 4]), + 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) def test_magnet_cylinder_segment_Hfield(): - Btest = np.array( + Btest = xp.asarray( [ [-1948.14367497, 32319.94437208, 17616.88571231], [14167.64961763, 1419.94126065, 17921.6463117], ] ) B = magnet_cylinder_segment_Hfield( - observers=np.array([(1, 1, 2), (0, 0, 0)]), - dimensions=np.array([(1, 2, 0.1, 0.2, -1, 1), (1, 2, 0.3, 0.9, 0, 1)]), - magnetizations=np.array([(1e7, 0.1, 0.2), (1e6, 1.1, 2.2)]), + 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) def test_triangle_Bfield(): - Btest = np.array( + Btest = xp.asarray( [[7.45158965, 4.61994866, 3.13614132], [2.21345618, 2.67710148, 2.21345618]] ) B = triangle_Bfield( - observers=np.array([(2, 1, 1), (2, 2, 2)]), - vertices=np.array( - [[(0, 0, 0), (0, 0, 1), (1, 0, 0)], [(0, 0, 0), (0, 0, 1), (1, 0, 0)]] + 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=np.array([(1, 1, 1), (1, 1, 0)]) * 1e3, + polarizations=xp.asarray([(1.0, 1.0, 1.0), (1.0, 1.0, 0.0)]) * 1e3, ) np.testing.assert_allclose(B, Btest) From 6ad2478dea37a71451c85cf3803397d8385cf7bb Mon Sep 17 00:00:00 2001 From: purepani Date: Tue, 20 May 2025 17:31:37 -0500 Subject: [PATCH 04/16] Adds utility functions for array api Taken directly from https://github.com/scipy/scipy/blob/7f75a948bd46d0453885dac81efded3aaa640907/scipy/_lib/_array_api.py --- src/magpylib/_src/array_api_utils.py | 194 +++++++++++++++++++++++++++ 1 file changed, 194 insertions(+) create mode 100644 src/magpylib/_src/array_api_utils.py diff --git a/src/magpylib/_src/array_api_utils.py b/src/magpylib/_src/array_api_utils.py new file mode 100644 index 000000000..89af5adff --- /dev/null +++ b/src/magpylib/_src/array_api_utils.py @@ -0,0 +1,194 @@ +import numpy as np +import numpy.typing as npt +from array_api_compat import ( + is_numpy_namespace as is_numpy, + is_torch_namespace as is_torch, + array_namespace, +) +from typing import Any, TypeAlias, Literal +from types import ModuleType + +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() + else: + # 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) From 3167d02feadff01034efed9cd96a7dc16c0767a3 Mon Sep 17 00:00:00 2001 From: purepani Date: Wed, 7 May 2025 19:06:33 -0500 Subject: [PATCH 05/16] Adds array api support for cuboid core function --- src/magpylib/_src/fields/field_BH_cuboid.py | 112 +++++++++++--------- 1 file changed, 59 insertions(+), 53 deletions(-) diff --git a/src/magpylib/_src/fields/field_BH_cuboid.py b/src/magpylib/_src/fields/field_BH_cuboid.py index d79a47b2f..7f3a82e02 100644 --- a/src/magpylib/_src/fields/field_BH_cuboid.py +++ b/src/magpylib/_src/fields/field_BH_cuboid.py @@ -6,9 +6,11 @@ from __future__ import annotations import numpy as np +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 # pylint: disable=too-many-statements @@ -79,9 +81,14 @@ def magnet_cuboid_Bfield( Cichon: IEEE Sensors Journal, vol. 19, no. 7, April 1, 2019, p.2509 """ - pol_x, pol_y, pol_z = polarizations.T - a, b, c = dimensions.T / 2 - x, y, z = np.copy(observers).T + xp = array_namespace(observers, dimensions, polarizations) + observers, dimensions, polarizations = xp_promote( + observers, dimensions, polarizations, force_floating=True, xp=xp + ) + + pol_x, pol_y, pol_z = (polarizations[:, i] for i in range(3)) + a, b, c = ((dimensions / 2)[:, i] for i in range(3)) + x, y, z = (observers[:, i] for i in range(3)) # avoid indeterminate forms by evaluating in bottQ4 only -------- # basic masks @@ -95,10 +102,10 @@ def magnet_cuboid_Bfield( z[maskz] = z[maskz] * -1 # create sign flips for position changes - qsigns = np.ones((len(pol_x), 3, 3)) - qs_flipx = np.array([[1, -1, -1], [-1, 1, 1], [-1, 1, 1]]) - qs_flipy = np.array([[1, -1, 1], [-1, 1, -1], [1, -1, 1]]) - qs_flipz = np.array([[1, 1, -1], [1, 1, -1], [-1, -1, 1]]) + qsigns = xp.ones((pol_x.shape[0], 3, 3), dtype=xp.float32) + qs_flipx = xp.asarray([[1, -1, -1], [-1, 1, 1], [-1, 1, 1]], dtype=xp.float32) + qs_flipy = xp.asarray([[1, -1, 1], [-1, 1, -1], [1, -1, 1]], dtype=xp.float32) + qs_flipz = xp.asarray([[1, 1, -1], [1, 1, -1], [-1, -1, 1]], dtype=xp.float32) # signs flips can be applied subsequently qsigns[maskx] = qsigns[maskx] * qs_flipx qsigns[masky] = qsigns[masky] * qs_flipy @@ -119,59 +126,58 @@ def magnet_cuboid_Bfield( ymb2, ypb2 = ymb**2, ypb**2 zmc2, zpc2 = zmc**2, zpc**2 - mmm = np.sqrt(xma2 + ymb2 + zmc2) - pmp = np.sqrt(xpa2 + ymb2 + zpc2) - pmm = np.sqrt(xpa2 + ymb2 + zmc2) - mmp = np.sqrt(xma2 + ymb2 + zpc2) - mpm = np.sqrt(xma2 + ypb2 + zmc2) - ppp = np.sqrt(xpa2 + ypb2 + zpc2) - ppm = np.sqrt(xpa2 + ypb2 + zmc2) - mpp = np.sqrt(xma2 + ypb2 + zpc2) - - with np.errstate(divide="ignore", invalid="ignore"): - ff2x = np.log((xma + mmm) * (xpa + ppm) * (xpa + pmp) * (xma + mpp)) - np.log( - (xpa + pmm) * (xma + mpm) * (xma + mmp) * (xpa + ppp) - ) + mmm = xp.sqrt(xma2 + ymb2 + zmc2) + pmp = xp.sqrt(xpa2 + ymb2 + zpc2) + pmm = xp.sqrt(xpa2 + ymb2 + zmc2) + mmp = xp.sqrt(xma2 + ymb2 + zpc2) + mpm = xp.sqrt(xma2 + ypb2 + zmc2) + ppp = xp.sqrt(xpa2 + ypb2 + zpc2) + ppm = xp.sqrt(xpa2 + ypb2 + zmc2) + mpp = xp.sqrt(xma2 + ypb2 + zpc2) + + ff2x = xp.log((xma + mmm) * (xpa + ppm) * (xpa + pmp) * (xma + mpp)) - xp.log( + (xpa + pmm) * (xma + mpm) * (xma + mmp) * (xpa + ppp) + ) - ff2y = np.log( - (-ymb + mmm) * (-ypb + ppm) * (-ymb + pmp) * (-ypb + mpp) - ) - np.log((-ymb + pmm) * (-ypb + mpm) * (ymb - mmp) * (ypb - ppp)) + ff2y = xp.log((-ymb + mmm) * (-ypb + ppm) * (-ymb + pmp) * (-ypb + mpp)) - xp.log( + (-ymb + pmm) * (-ypb + mpm) * (ymb - mmp) * (ypb - ppp) + ) - ff2z = np.log( - (-zmc + mmm) * (-zmc + ppm) * (-zpc + pmp) * (-zpc + mpp) - ) - np.log((-zmc + pmm) * (zmc - mpm) * (-zpc + mmp) * (zpc - ppp)) + ff2z = xp.log((-zmc + mmm) * (-zmc + ppm) * (-zpc + pmp) * (-zpc + mpp)) - xp.log( + (-zmc + pmm) * (zmc - mpm) * (-zpc + mmp) * (zpc - ppp) + ) ff1x = ( - np.arctan2((ymb * zmc), (xma * mmm)) - - np.arctan2((ymb * zmc), (xpa * pmm)) - - np.arctan2((ypb * zmc), (xma * mpm)) - + np.arctan2((ypb * zmc), (xpa * ppm)) - - np.arctan2((ymb * zpc), (xma * mmp)) - + np.arctan2((ymb * zpc), (xpa * pmp)) - + np.arctan2((ypb * zpc), (xma * mpp)) - - np.arctan2((ypb * zpc), (xpa * ppp)) + xp.atan2((ymb * zmc), (xma * mmm)) + - xp.atan2((ymb * zmc), (xpa * pmm)) + - xp.atan2((ypb * zmc), (xma * mpm)) + + xp.atan2((ypb * zmc), (xpa * ppm)) + - xp.atan2((ymb * zpc), (xma * mmp)) + + xp.atan2((ymb * zpc), (xpa * pmp)) + + xp.atan2((ypb * zpc), (xma * mpp)) + - xp.atan2((ypb * zpc), (xpa * ppp)) ) ff1y = ( - np.arctan2((xma * zmc), (ymb * mmm)) - - np.arctan2((xpa * zmc), (ymb * pmm)) - - np.arctan2((xma * zmc), (ypb * mpm)) - + np.arctan2((xpa * zmc), (ypb * ppm)) - - np.arctan2((xma * zpc), (ymb * mmp)) - + np.arctan2((xpa * zpc), (ymb * pmp)) - + np.arctan2((xma * zpc), (ypb * mpp)) - - np.arctan2((xpa * zpc), (ypb * ppp)) + xp.atan2((xma * zmc), (ymb * mmm)) + - xp.atan2((xpa * zmc), (ymb * pmm)) + - xp.atan2((xma * zmc), (ypb * mpm)) + + xp.atan2((xpa * zmc), (ypb * ppm)) + - xp.atan2((xma * zpc), (ymb * mmp)) + + xp.atan2((xpa * zpc), (ymb * pmp)) + + xp.atan2((xma * zpc), (ypb * mpp)) + - xp.atan2((xpa * zpc), (ypb * ppp)) ) ff1z = ( - np.arctan2((xma * ymb), (zmc * mmm)) - - np.arctan2((xpa * ymb), (zmc * pmm)) - - np.arctan2((xma * ypb), (zmc * mpm)) - + np.arctan2((xpa * ypb), (zmc * ppm)) - - np.arctan2((xma * ymb), (zpc * mmp)) - + np.arctan2((xpa * ymb), (zpc * pmp)) - + np.arctan2((xma * ypb), (zpc * mpp)) - - np.arctan2((xpa * ypb), (zpc * ppp)) + xp.atan2((xma * ymb), (zmc * mmm)) + - xp.atan2((xpa * ymb), (zmc * pmm)) + - xp.atan2((xma * ypb), (zmc * mpm)) + + xp.atan2((xpa * ypb), (zmc * ppm)) + - xp.atan2((xma * ymb), (zpc * mmp)) + + xp.atan2((xpa * ymb), (zpc * pmp)) + + xp.atan2((xma * ypb), (zpc * mpp)) + - xp.atan2((xpa * ypb), (zpc * ppp)) ) # contributions from x-polarization @@ -193,10 +199,10 @@ def magnet_cuboid_Bfield( by_tot = by_pol_x + by_pol_y + by_pol_z bz_tot = bz_pol_x + bz_pol_y + bz_pol_z - # B = np.c_[bx_tot, by_tot, bz_tot] # faster for 10^5 and more evaluations - B = np.concatenate(((bx_tot,), (by_tot,), (bz_tot,)), axis=0).T + # B = xp.c_[bx_tot, by_tot, bz_tot] # faster for 10^5 and more evaluations + B = xp.stack((bx_tot, by_tot, bz_tot), axis=-1) - B /= 4 * np.pi + B /= 4 * xp.pi return B From b89c186a9d1341ae174b6b42f1d78dca138d9cf7 Mon Sep 17 00:00:00 2001 From: purepani Date: Wed, 7 May 2025 23:49:59 -0500 Subject: [PATCH 06/16] Adds array api support for triangle core function --- src/magpylib/_src/fields/field_BH_triangle.py | 143 +++++++++++++----- 1 file changed, 101 insertions(+), 42 deletions(-) diff --git a/src/magpylib/_src/fields/field_BH_triangle.py b/src/magpylib/_src/fields/field_BH_triangle.py index c311a330b..925d06068 100644 --- a/src/magpylib/_src/fields/field_BH_triangle.py +++ b/src/magpylib/_src/fields/field_BH_triangle.py @@ -6,13 +6,18 @@ # pylance: disable=Code is unreachable from __future__ import annotations +from functools import partial + +import array_api_extra as xpx import numpy as np +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 -def vcross3(a: np.ndarray, b: np.ndarray) -> np.ndarray: +def vcross3(xp, a: np.ndarray, b: np.ndarray) -> np.ndarray: """ vectorized cross product for 3d vectors. Is ~4x faster than np.cross when arrays are smallish. Only slightly faster for large arrays. @@ -20,29 +25,28 @@ def vcross3(a: np.ndarray, b: np.ndarray) -> np.ndarray: returns: (n, 3) """ # receives nan values at corners - with np.errstate(invalid="ignore"): - result = np.array( - [ - a[:, 1] * b[:, 2] - a[:, 2] * b[:, 1], - a[:, 2] * b[:, 0] - a[:, 0] * b[:, 2], - a[:, 0] * b[:, 1] - a[:, 1] * b[:, 0], - ] - ) + result = xp.asarray( + [ + a[:, 1] * b[:, 2] - a[:, 2] * b[:, 1], + a[:, 2] * b[:, 0] - a[:, 0] * b[:, 2], + a[:, 0] * b[:, 1] - a[:, 1] * b[:, 0], + ] + ) return result.T -def norm_vector(v) -> np.ndarray: +def norm_vector_xp(xp, v) -> np.ndarray: """ Calculates normalized orthogonal vector on a plane defined by three vertices. """ - a = v[:, 1] - v[:, 0] - b = v[:, 2] - v[:, 0] - n = vcross3(a, b) - n_norm = np.linalg.norm(n, axis=-1) - return n / np.expand_dims(n_norm, axis=-1) + a = v[:, 1, ...] - v[:, 0, ...] + b = v[:, 2, ...] - v[:, 0, ...] + n = xp.linalg.cross(a, b) + n_norm = xp.sqrt(xp.vecdot(n, n, axis=-1)) + return n / xp.expand_dims(n_norm, axis=-1) -def solid_angle(R: np.ndarray, r: np.ndarray) -> np.ndarray: +def solid_angle_xp(xp, R: np.ndarray, r: np.ndarray) -> np.ndarray: """ Vectorized computation of the solid angle of triangles. @@ -56,23 +60,40 @@ def solid_angle(R: np.ndarray, r: np.ndarray) -> np.ndarray: Returns: [sangle_a, sangle_b, sangle_c, ...] """ + R0, R1, R2 = R[0, ...], R[1, ...], R[2, ...] + r0, r1, r2 = r[0, ...], r[1, ...], r[2, ...] # Calculates (oriented) volume of the parallelepiped in vectorized form. - N = np.einsum("ij, ij->i", R[2], vcross3(R[1], R[0])) + N = xp.vecdot( + R2, + xp.linalg.cross(R1, R0), + ) D = ( - r[0] * r[1] * r[2] - + np.einsum("ij, ij->i", R[2], R[1]) * r[0] - + np.einsum("ij, ij->i", R[2], R[0]) * r[1] - + np.einsum("ij, ij->i", R[1], R[0]) * r[2] + r0 * r1 * r2 + + xp.vecdot( + R2, + R1, + ) + * r0 + + xp.vecdot( + R2, + R0, + ) + * r1 + + xp.vecdot( + R1, + R0, + ) + * r2 ) - result = 2.0 * np.arctan2(N, D) + result = 2.0 * xp.atan2(N, D) # modulus 2pi to avoid jumps on edges in line # "B = sigma * ((n.T * solid_angle(R, r)) - vcross3(n, PQR).T)" # <-- bad fix :( - return np.where(abs(result) > 6.2831853, 0, result) + return xp.where(xp.abs(result) > 6.2831853, 0, result) def triangle_Bfield( @@ -130,24 +151,37 @@ def triangle_Bfield( Loss of precision when approaching a triangle as (x-edge)**2 :( Loss of precision with distance from the triangle as distance**3 :( """ + xp = array_namespace(observers, vertices, polarizations) + observers, vertices, polarizations = xp_promote( + observers, vertices, polarizations, force_floating=True, xp=xp + ) + norm_vector = partial(norm_vector_xp, xp) + solid_angle = partial(solid_angle_xp, xp) + n = norm_vector(vertices) - sigma = np.einsum("ij, ij->i", n, polarizations) # vectorized inner product + sigma = xp.vecdot( + n, + polarizations, + ) # vectorized inner product # vertex <-> observer - R = np.swapaxes(vertices, 0, 1) - observers - r2 = np.sum(R * R, axis=-1) - r = np.sqrt(r2) + R = xp.permute_dims(vertices, (1, 0, 2)) - observers + r2 = xp.sum(R * R, axis=-1) + r = xp.sqrt(r2) # vertex <-> vertex - L = vertices[:, (1, 2, 0)] - vertices[:, (0, 1, 2)] - L = np.swapaxes(L, 0, 1) - l2 = np.sum(L * L, axis=-1) - l1 = np.sqrt(l2) + L = xp.empty_like(vertices[:, 0:3, :]) + L = xpx.at(L)[:, 0, ...].set(vertices[:, 1, :] - vertices[:, 0, :]) + L = xpx.at(L)[:, 1, ...].set(vertices[:, 2, :] - vertices[:, 1, :]) + L = xpx.at(L)[:, 2, ...].set(vertices[:, 0, :] - vertices[:, 2, :]) + L = xp.permute_dims(L, (1, 0, 2)) + l2 = xp.sum(L * L, axis=-1) + l1 = xp.sqrt(l2) # vert-vert -- vert-obs - b = np.einsum("ijk, ijk->ij", R, L) + b = xp.vecdot(R, L) bl = b / l1 - ind = np.fabs(r + bl) # closeness measure to corner and edge + ind = xp.abs(r + bl) # closeness measure to corner and edge # The computation of ind is the origin of a major numerical instability # when approaching the triangle because r ~ -bl. This number @@ -166,16 +200,41 @@ def triangle_Bfield( # 1/x, # 0 # ) + def B_not_degenerate(l1, l2, r, r2, b, bl, ind): + l_arg1 = xp.sqrt(l2 + 2 * b + r2) + l1 + bl + return 1.0 / l1 * xp.log(l_arg1 / ind) + + def B_edge(l1, r): + lr = xp.abs(l1 - r) + return -(1.0 / l1) * xp.log(lr / r) + + def B_degenerate(l1, l2, r, r2, b, bl, ind): + r_mask = (r != 0.0) & (xp.abs(l1 - r) > 0.0) + return xpx.apply_where(r_mask, (l1, r), B_edge, fill_value=xp.nan) + + ind_mask = ind > 1.0e-12 + r_mask = r == 0.0 + r_copy = xp.asarray(r, copy=True) + lr = xp.abs(l1 - r) + lr_mask = lr == 0.0 + lr = xpx.at(lr)[ind_mask | lr_mask].set(3.0) + ind = xpx.at(ind)[~ind_mask].set(1.0) + r_copy = xpx.at(r_copy)[r_mask].set(1.0) + + I = xpx.apply_where( # noqa: E741 + ind_mask, + (l1, l2, r, r2, b, bl, ind), + B_not_degenerate, + B_degenerate, + # 1.0 / l1 * xp.log(l_arg1 / ind), + # -(1.0 / l1) * xp.log(lr / r_copy), + ) + I = xpx.at(I)[r_mask].set(xp.nan) + I = xpx.at(I)[lr_mask & ~ind_mask].set(xp.nan) - with np.errstate(divide="ignore", invalid="ignore"): - I = np.where( # noqa: E741 - ind > 1.0e-12, - 1.0 / l1 * np.log((np.sqrt(l2 + 2 * b + r2) + l1 + bl) / ind), - -(1.0 / l1) * np.log(np.fabs(l1 - r) / r), - ) - PQR = np.einsum("ij, ijk -> jk", I, L) - B = sigma * (n.T * solid_angle(R, r) - vcross3(n, PQR).T) - B = B / np.pi / 4.0 + PQR = xp.vecdot(I[:, :, xp.newaxis], L, axis=-3) + B = sigma * (n.T * solid_angle(R, r) - xp.linalg.cross(n, PQR).T) + B = B / xp.pi / 4.0 return B.T From 10be96afbffe35e3d02589e9256a642a95e07747 Mon Sep 17 00:00:00 2001 From: purepani Date: Wed, 7 May 2025 20:46:56 -0500 Subject: [PATCH 07/16] Adds array api support for polyline core function --- src/magpylib/_src/fields/field_BH_polyline.py | 71 ++++++++++--------- 1 file changed, 38 insertions(+), 33 deletions(-) diff --git a/src/magpylib/_src/fields/field_BH_polyline.py b/src/magpylib/_src/fields/field_BH_polyline.py index 7d998a321..21fadafd6 100644 --- a/src/magpylib/_src/fields/field_BH_polyline.py +++ b/src/magpylib/_src/fields/field_BH_polyline.py @@ -5,11 +5,15 @@ # pylint: disable=too-many-positional-arguments from __future__ import annotations +from itertools import zip_longest + +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.input_checks import check_field_input +from magpylib._src.array_api_utils import xp_promote def current_vertices_field( @@ -124,72 +128,73 @@ 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) # 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] + 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) + cros_ = xp.linalg.cross(p2 - p1, po - p4) + + norm_cros = xp.linalg.vector_norm(cros_, axis=-1) eB = (cros_.T / norm_cros).T # compute angles - norm_o1 = norm( - po - p1, axis=1 + norm_o1 = xp.linalg.vector_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) + 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 - deltaSin = np.empty((len(po),)) + 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[mask2] = abs(sinTh1[mask2] - sinTh2[mask2]) + 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[mask3] = abs(sinTh2[mask3] - sinTh1[mask3]) + 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[mask4] = abs(sinTh1[mask4] + sinTh2[mask4]) + 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 # 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 + H = xp.zeros_like(observers, dtype=xp.float64) + H = xpx.at(H)[~mask1].set( + (deltaSin / norm_o4 * eB.T / norm_12 * currents / (4 * np.pi)).T + ) + return H def BHJM_current_polyline( From 6347cc489fca6ebe0e4f9f22e7243b998e528f7c Mon Sep 17 00:00:00 2001 From: purepani Date: Tue, 6 May 2025 11:12:13 -0500 Subject: [PATCH 08/16] Adds array api support for dipole field core function --- src/magpylib/_src/fields/field_BH_dipole.py | 41 ++++++++++----------- 1 file changed, 20 insertions(+), 21 deletions(-) diff --git a/src/magpylib/_src/fields/field_BH_dipole.py b/src/magpylib/_src/fields/field_BH_dipole.py index 07196bc76..201f008b0 100644 --- a/src/magpylib/_src/fields/field_BH_dipole.py +++ b/src/magpylib/_src/fields/field_BH_dipole.py @@ -4,10 +4,13 @@ from __future__ import annotations +import array_api_extra as xpx import numpy as np +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 # CORE @@ -53,27 +56,23 @@ def dipole_Hfield( ----- The moment of a magnet is given by its volume*magnetization. """ - - x, y, z = observers.T - r = np.sqrt(x**2 + y**2 + z**2) # faster than np.linalg.norm - with np.errstate(divide="ignore", invalid="ignore"): - # 0/0 produces invalid warn and results in np.nan - # x/0 produces divide warn and results in np.inf - H = ( - ( - 3 * np.sum(moments * observers, axis=1) * observers.T / r**5 - - moments.T / r**3 - ).T - / 4 - / np.pi - ) - - # when r=0 return np.inf in all non-zero moments directions - mask1 = r == 0 - if np.any(mask1): - with np.errstate(divide="ignore", invalid="ignore"): - H[mask1] = moments[mask1] / 0.0 - np.nan_to_num(H, copy=False, posinf=np.inf, neginf=-np.inf) + xp = array_namespace(observers, moments) + observers, moments = xp_promote(observers, moments, force_floating=True, xp=xp) + r = xp.linalg.vector_norm(observers, axis=-1, keepdims=True) + + # 0/0 produces invalid warn and results in np.nan + # x/0 produces divide warn and results in np.inf + mask_r = r == 0.0 + r = xpx.at(r)[mask_r].set(1.0) + dotprod = xp.vecdot(moments, observers)[:, xp.newaxis] + + def B(dotprod, observers, moments, r): + A = xp.divide(3 * dotprod * observers, r**5) + B = xp.divide(moments, r**3) + return xp.divide((A - B), 4.0 * xp.pi) + + H = xpx.apply_where(~mask_r, (dotprod, observers, moments, r), B, fill_value=xp.inf) + # H = xpx.at(H)[xp.broadcast_to(mask_r, H.shape)].set(xp.inf) return H From 41df3314d6dc9975cd8459f5885198b2e62655ab Mon Sep 17 00:00:00 2001 From: purepani Date: Tue, 6 May 2025 18:35:38 -0500 Subject: [PATCH 09/16] Adds array api support for sphere core function --- src/magpylib/_src/fields/field_BH_sphere.py | 42 ++++++++++++++++++--- 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/src/magpylib/_src/fields/field_BH_sphere.py b/src/magpylib/_src/fields/field_BH_sphere.py index 585ba1e2d..a9c64c446 100644 --- a/src/magpylib/_src/fields/field_BH_sphere.py +++ b/src/magpylib/_src/fields/field_BH_sphere.py @@ -5,10 +5,13 @@ from __future__ import annotations +import array_api_extra as xpx import numpy as np +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 # CORE @@ -58,12 +61,39 @@ def magnet_sphere_Bfield( The field corresponds to a dipole field on the outside and is 2/3*mag in the inside (see e.g. "Theoretical Physics, Bertelmann"). """ - return BHJM_magnet_sphere( - field="B", - observers=observers, - diameter=diameters, - polarization=polarizations, - ) + xp = array_namespace(observers, diameters, polarizations) + 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 + + # inside field & allocate + polarization = polarizations + mask_out = r > r_sphere + + mask_in = ~mask_out + + def B_in(polarization): + return 2.0 * polarization / 3.0 + + def B_out(polarization, observers, r_sphere, r): + mask = r == 0.0 + r = xpx.at(r)[mask].set(xp.nan) + return ( + ( + xp.divide( + 3 * xp.sum(polarization * observers, axis=1) * observers.T + - polarization.T * r**2, + r**5, + ) + ) + * r_sphere**3 + / 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])) + + return B def BHJM_magnet_sphere( From 07054b37337f79d7984b5b354f754700c25454e6 Mon Sep 17 00:00:00 2001 From: purepani Date: Mon, 19 May 2025 14:00:10 -0500 Subject: [PATCH 10/16] Adds array api support for circle core function --- src/magpylib/_src/fields/field_BH_circle.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/magpylib/_src/fields/field_BH_circle.py b/src/magpylib/_src/fields/field_BH_circle.py index 9ee4fc062..46c87e32d 100644 --- a/src/magpylib/_src/fields/field_BH_circle.py +++ b/src/magpylib/_src/fields/field_BH_circle.py @@ -8,9 +8,12 @@ import numpy as np from scipy.constants import mu_0 as MU0 +from array_api_compat import array_namespace + from magpylib._src.fields.special_cel import cel_iter from magpylib._src.input_checks import check_field_input from magpylib._src.utility import cart_to_cyl_coordinates, cyl_field_to_cart +from magpylib._src.array_api_utils import xp_promote # CORE @@ -67,7 +70,9 @@ def current_circle_Hfield( efficient expression for the magnetic field of a current loop.", M.Ortner et al, Magnetism 2023, 3(1), 11-31. """ - n5 = len(r) + xp = array_namespace(r0, r, z, i0) + r0, r, z, i0 = xp_promote(r0, r, z, i0, force_floating=True, xp=xp) + n5 = r.shape[0] # express through ratios (make dimensionless, avoid large/small input values, stupid) r = r / r0 @@ -79,23 +84,23 @@ def current_circle_Hfield( k2 = 4 * r / x0 q2 = (z2 + (r - 1) ** 2) / x0 - k = np.sqrt(k2) - q = np.sqrt(q2) + k = xp.sqrt(k2) + q = xp.sqrt(q2) p = 1 + q - pf = k / np.sqrt(r) / q2 / 20 / r0 * 1e-6 * i0 + pf = k / xp.sqrt(r) / q2 / 20 / r0 * 1e-6 * i0 # cel* part cc = k2 * k2 ss = 2 * cc * q / p - Hr = pf * z / r * cel_iter(q, p, np.ones(n5), cc, ss, p, q) + Hr = pf * z / r * cel_iter(q, p, xp.ones(n5), cc, ss, p, q) # cel** part cc = k2 * (k2 - (q2 + 1) / r) ss = 2 * k2 * q * (k2 / p - p / r) - Hz = -pf * cel_iter(q, p, np.ones(n5), cc, ss, p, q) + Hz = -pf * cel_iter(q, p, xp.ones(n5), cc, ss, p, q) # input is I -> output must be H-field - return np.vstack((Hr, np.zeros(n5), Hz)) * 795774.7154594767 # *1e7/4/np.pi + return xp.stack((Hr, xp.zeros(n5), Hz), axis=0) * 795774.7154594767 # *1e7/4/np.pi def BHJM_circle( From de2950cfa1ee8a6c8a5d260a033388a8dc371a3c Mon Sep 17 00:00:00 2001 From: purepani Date: Tue, 20 May 2025 16:26:39 -0500 Subject: [PATCH 11/16] Adds tests for multiple backends(numpy, jax, array_api_strict, dask). --- pyproject.toml | 4 +++ tests/test_core.py | 62 ++++++++++++++++++++++++++++++---------------- 2 files changed, 45 insertions(+), 21 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3b6f44814..e001243fe 100644 --- 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/tests/test_core.py b/tests/test_core.py index d4585f740..d9408d7b8 100644 --- 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,8 +20,24 @@ 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)]), @@ -25,10 +45,10 @@ def test_magnet_sphere_Bfield(): 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]), @@ -36,10 +56,10 @@ def test_current_circle_Hfield(): 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,14 +82,14 @@ 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( @@ -77,20 +97,20 @@ def test_magnet_cuboid_Bfield(): 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) From bcae406746e23dfac664252fc05b579831dd4c8c Mon Sep 17 00:00:00 2001 From: purepani Date: Sun, 1 Jun 2025 19:36:29 -0500 Subject: [PATCH 12/16] Implements elliptic functions for array api --- src/magpylib/_src/fields/special_cel.py | 108 ++++---- src/magpylib/_src/fields/special_el3.py | 269 ++++++++++--------- src/magpylib/_src/fields/special_elliptic.py | 24 ++ 3 files changed, 225 insertions(+), 176 deletions(-) create mode 100644 src/magpylib/_src/fields/special_elliptic.py diff --git a/src/magpylib/_src/fields/special_cel.py b/src/magpylib/_src/fields/special_cel.py index 149dc6880..715c5f001 100644 --- a/src/magpylib/_src/fields/special_cel.py +++ b/src/magpylib/_src/fields/special_cel.py @@ -3,7 +3,9 @@ import math as m +import array_api_extra as xpx import numpy as np +from array_api_compat import array_namespace def cel0(kc, p, c, s): @@ -56,63 +58,74 @@ def celv(kc, p, c, s): """ vectorized version of the cel integral above """ + xp = array_namespace(kc, p, c, s) # if kc == 0: # return NaN + # errtol = 0.000001 errtol = 0.000001 - n = len(kc) + n = kc.shape[0] - k = np.abs(kc) - em = np.ones(n, dtype=float) + k = xp.abs(kc) + em = xp.ones(n, dtype=xp.float64) - cc = c.copy() - pp = p.copy() - ss = s.copy() + # cc = xp.asarray(c, copy=True) + # pp = xp.asarray(p, copy=True) + # ss = xp.asarray(s, copy=True) # apply a mask for evaluation of respective cases mask = p <= 0 - # if p>0: - pp[~mask] = np.sqrt(p[~mask]) - ss[~mask] = s[~mask] / pp[~mask] - - # else: - f = kc[mask] * kc[mask] + f = kc * kc q = 1.0 - f - g = 1.0 - pp[mask] - f = f - pp[mask] - q = q * (ss[mask] - c[mask] * pp[mask]) - pp[mask] = np.sqrt(f / g) - cc[mask] = (c[mask] - ss[mask]) / g - ss[mask] = -q / (g * g * pp[mask]) + cc[mask] * pp[mask] - - f = cc.copy() + g = 1.0 - p + f = f - p + q = q * (s - c * p) + pp = xpx.apply_where( + mask, + (p, f, g), + lambda p, f, g: xp.sqrt(f / g), + lambda p, f, g: xp.sqrt(p), + ) + cc = xpx.apply_where( + mask, (c, s, g), lambda c, s, g: (c - s) / g, lambda c, ss, g: c + ) + ss = xpx.apply_where( + mask, + (s, pp, q, g, cc), + lambda s, pp, q, g, cc: -q / (g * g * pp) + cc * pp, + lambda s, pp, q, g, cc: s / pp, + ) + + f = xp.asarray(cc, copy=True) cc = cc + ss / pp g = k / pp ss = 2 * (ss + f * g) pp = g + pp - g = em.copy() + g = xp.asarray(em, copy=True) em = k + em - kk = k.copy() + kk = xp.asarray(k, copy=True) # define a mask that adjusts with every evaluation step so that only # non-converged entries are further iterated. - mask = np.ones(n, dtype=bool) - while np.any(mask): - k[mask] = 2 * np.sqrt(kk[mask]) - kk[mask] = k[mask] * em[mask] - f[mask] = cc[mask] - cc[mask] = cc[mask] + ss[mask] / pp[mask] - g[mask] = kk[mask] / pp[mask] - ss[mask] = 2 * (ss[mask] + f[mask] * g[mask]) - pp[mask] = g[mask] + pp[mask] - g[mask] = em[mask] - em[mask] = k[mask] + em[mask] + mask = xp.ones(n, dtype=xp.bool) + while xp.any(mask): + k = 2 * xp.sqrt(kk) + kk = k * em + f = cc + cc = cc + ss / pp + g = kk / pp + ss = 2 * (ss + f * g) + pp = g + pp + g = em + em = k + em # redefine mask - mask = np.abs(g - k) > g * errtol + err = g - k + tol = g * errtol + mask = xp.abs(err) > tol - return (np.pi / 2) * (ss + cc * em) / (em * (em + pp)) + return (xp.pi / 2) * (ss + cc * em) / (em * (em + pp)) def cel(kcv: np.ndarray, pv: np.ndarray, cv: np.ndarray, sv: np.ndarray) -> np.ndarray: @@ -128,12 +141,12 @@ def ellipticE(x): def ellipticPi(x, y): return elliptic((1-y)**(1/2.), 1-x, 1, 1) """ - n_input = len(kcv) + # n_input = len(kcv) - if n_input < 10: - return np.array( - [cel0(kc, p, c, s) for kc, p, c, s in zip(kcv, pv, cv, sv, strict=False)] - ) + # if n_input < 10: + # return np.array( + # [cel0(kc, p, c, s) for kc, p, c, s in zip(kcv, pv, cv, sv, strict=False)] + # ) return celv(kcv, pv, cv, sv) @@ -146,11 +159,11 @@ def cel_iter(qc, p, g, cc, ss, em, kk): # This cannot happen in core functions # case2: small input vector - loop is faster than vectorized computation - n_input = len(qc) - if n_input < 15: - result = np.zeros(n_input) - for i in range(n_input): - result[i] = cel_iter0(qc[i], p[i], g[i], cc[i], ss[i], em[i], kk[i]) + # n_input = len(qc) + # if n_input < 15: + # result = np.zeros(n_input) + # for i in range(n_input): + # result[i] = cel_iter0(qc[i], p[i], g[i], cc[i], ss[i], em[i], kk[i]) # case3: vectorized evaluation return cel_iterv(qc, p, g, cc, ss, em, kk) @@ -177,8 +190,9 @@ def cel_iterv(qc, p, g, cc, ss, em, kk): """ Iterative part of Bulirsch cel algorithm """ - while np.any(np.fabs(g - qc) >= qc * 1e-8): - qc = 2 * np.sqrt(kk) + xp = array_namespace(qc, p, g, cc, ss, em, kk) + while xp.any(xp.abs(g - qc) >= qc * 1e-8): + qc = 2 * xp.sqrt(kk) kk = qc * em f = cc cc = cc + ss / p diff --git a/src/magpylib/_src/fields/special_el3.py b/src/magpylib/_src/fields/special_el3.py index e2578db9c..8d1d35298 100644 --- a/src/magpylib/_src/fields/special_el3.py +++ b/src/magpylib/_src/fields/special_el3.py @@ -1,6 +1,7 @@ from __future__ import annotations import numpy as np +from array_api_compat import array_namespace from magpylib._src.fields.special_cel import cel @@ -64,7 +65,7 @@ def el30(x, kc, p): if bo: s = -s u = (u + 1) * 0.5 - return (u - s * h) * np.sqrt(h) * x + u * np.arcsinh(x) + return (u - s * h) * np.sqrt(h) * x + u * np.asinh(x) w = 1 + f if w == 0: @@ -217,7 +218,7 @@ def el30(x, kc, p): break if y < 0.0: l += 1 - e = np.arctan(t / y) + np.pi * l + e = np.atan(t / y) + np.pi * l e = e * (c * t + d) / (t * (t + q)) if bo: h = v / (t + u) @@ -227,9 +228,9 @@ def el30(x, kc, p): z = CB if z < 0.0: m = m + np.sign(h) - s = np.arctan(h / z) + m * np.pi + s = np.atan(h / z) + m * np.pi else: - s = np.arcsinh(ye) if bk else np.log(z) + m * ln2 + s = np.asinh(ye) if bk else np.log(z) + m * ln2 s = s * 0.5 e = (e + np.sqrt(fa) * s) / n return e if (x > 0.0) else -e @@ -249,8 +250,9 @@ def el3v(x, kc, p): # pylint: disable=too-many-branches # pylint: disable=too-many-statements - nnn0 = len(x) - result0 = np.zeros(nnn0) + xp = array_namespace(x, kc, p) + nnn0 = x.shape[0] + result0 = xp.zeros(nnn0) # return 0 when mask0 mask0 = x != 0 @@ -258,30 +260,31 @@ def el3v(x, kc, p): kc = kc[mask0] p = p[mask0] - nnn = len(x) - result = np.zeros(nnn) + nnn = x.shape[0] + result = xp.zeros(nnn) D = 8 CA = 10.0 ** (-D / 2) CB = 10.0 ** (-D - 2) ND = D - 2 - ln2 = np.log(2) + ln2 = xp.log(xp.asarray(2.0)) hh = x * x f = p * hh - s = np.zeros(nnn) + s = xp.zeros(nnn) mask1 = kc == 0 - s[mask1] = CA / (1 + np.abs(x[mask1])) + s[mask1] = CA / (1 + xp.abs(x[mask1])) s[~mask1] = kc[~mask1] t = s * s pm = 0.5 * t e = hh * t - z = np.abs(f) - r = np.abs(p) + z = xp.abs(f) + r = xp.abs(p) h = 1.0 + hh - mask2 = (e < 0.1) * (z < 0.1) * (t < 1) * (r < 1) - if any(mask2): - ra, rb, rr = np.zeros((3, ND - 1, np.sum(mask2))) + mask2 = (e < 0.1) & (z < 0.1) & (t < 1) & (r < 1) + if xp.any(mask2): + R = xp.zeros((3, ND - 1, xp.count_nonzero(mask2))) + ra, rb, rr = (R[i] for i in range(3)) px, pmx, tx, hhx, hx, xx = ( p[mask2], pm[mask2], @@ -302,7 +305,7 @@ def el3v(x, kc, p): pmx = pmx * tx * ra[km2] sx = sx * px + pmx ux = sx * zd - sx = np.copy(ux) + sx = xp.asarray(ux, copy=True) bo = False for k in range(ND, 1, -1): km2 = int(k - 2) @@ -313,7 +316,7 @@ def el3v(x, kc, p): if bo: sx = -sx ux = (ux + 1) * 0.5 - result[mask2] = (ux - sx * hx) * np.sqrt(hx) * xx + ux * np.arcsinh(xx) + result[mask2] = (ux - sx * hx) * xp.sqrt(hx) * xx + ux * xp.asinh(xx) mask2x = ~mask2 p, pm, t, hh, h, x = ( @@ -325,19 +328,20 @@ def el3v(x, kc, p): x[mask2x], ) f, e, z, s = f[mask2x], e[mask2x], z[mask2x], s[mask2x] - ye, k = np.zeros((2, len(p))) - bk = np.zeros(len(p), dtype=bool) + ye_k = xp.zeros((2, (p.shape[0]))) + ye, k = (ye_k[i, ...] for i in range(2)) + bk = xp.zeros((p.shape[0]), dtype=xp.bool) w = 1 + f - if np.any(w == 0): + if xp.any(w == 0): msg = "FAIL" raise RuntimeError(msg) - p1 = np.copy(p) + p1 = xp.asarray(p, copy=True) mask3 = p == 0 p1[mask3] = CB / hh[mask3] - s = np.abs(s) - y = np.abs(x) + s = xp.abs(s) + y = xp.abs(x) g = p1 - 1.0 g[g == 0] = CB f = p1 - t @@ -348,26 +352,27 @@ def el3v(x, kc, p): r = p1 * h fa = g / (f * p1) bo = fa > 0.0 - fa = np.abs(fa) - pz = np.abs(g * f) - de = np.sqrt(pz) - q = np.sqrt(np.abs(p1)) + fa = xp.abs(fa) + pz = xp.abs(g * f) + de = xp.sqrt(pz) + q = xp.sqrt(xp.abs(p1)) mask5 = pm > 0.5 pm[mask5] = 0.5 pm = p1 - pm - u, v, d, c = np.zeros((4, len(pm))) + uvdc = xp.zeros((4, (pm.shape[0]))) + u, v, d, c = (uvdc[i, ...] for i in range(4)) mask6 = pm >= 0.0 - if np.any(mask6): - u[mask6] = np.sqrt(r[mask6] * ap[mask6]) - v[mask6] = y[mask6] * de[mask6] * np.sign(g[mask6]) + if xp.any(mask6): + u[mask6] = xp.sqrt(r[mask6] * ap[mask6]) + v[mask6] = y[mask6] * de[mask6] * xp.sign(g[mask6]) d[mask6] = 1 / q[mask6] c[mask6] = 1.0 mask6x = ~mask6 - if np.any(mask6x): - u[mask6x] = np.sqrt(h[mask6x] * ap[mask6x] * pz[mask6x]) + if xp.any(mask6x): + u[mask6x] = xp.sqrt(h[mask6x] * ap[mask6x] * pz[mask6x]) ye[mask6x] = y[mask6x] * q[mask6x] v[mask6x] = am[mask6x] * ye[mask6x] q[mask6x] = -de[mask6x] / g[mask6x] @@ -375,74 +380,78 @@ def el3v(x, kc, p): c[mask6x] = 0 pz[mask6x] = ap[mask6x] - r[mask6x] - if np.any(bo): + if xp.any(bo): r[bo] = v[bo] / u[bo] z[bo] = 1 k[bo] = 1 - mask7 = bo * (pm < 0) - if np.any(mask7): - h[mask7] = y[mask7] * np.sqrt(h[mask7] / (ap[mask7] * fa[mask7])) + mask7 = bo & (pm < 0) + if xp.any(mask7): + h[mask7] = y[mask7] * xp.sqrt(h[mask7] / (ap[mask7] * fa[mask7])) h[mask7] = 1 / h[mask7] - h[mask7] z[mask7] = h[mask7] - 2 * r[mask7] r[mask7] = 2 + r[mask7] * h[mask7] - mask7a = mask7 * (r == 0) + mask7a = mask7 & (r == 0) r[mask7a] = CB - mask7b = mask7 * (z == 0) + mask7b = mask7 & (z == 0) z[mask7b] = h[mask7b] * CB z[mask7] = r[mask7] / z[mask7] - r[mask7] = np.copy(z[mask7]) + r[mask7] = xp.asarray(z[mask7], copy=True) w[mask7] = pz[mask7] u[bo] = u[bo] / w[bo] v[bo] = v[bo] / w[bo] box = ~bo - if np.any(box): - t[box] = u[box] + np.abs(v[box]) + if xp.any(box): + t[box] = u[box] + xp.abs(v[box]) bk[box] = True - mask8 = box * (p1 < 0) - if np.any(mask8): + mask8 = box & (p1 < 0) + if xp.any(mask8): de[mask8] = v[mask8] / pz[mask8] ye[mask8] = u[mask8] * ye[mask8] ye[mask8] = 2 * ye[mask8] u[mask8] = t[mask8] / pz[mask8] v[mask8] = (-f[mask8] - g[mask8] * e[mask8]) / t[mask8] - t[mask8] = pz[mask8] * np.abs(w[mask8]) + t[mask8] = pz[mask8] * xp.abs(w[mask8]) z[mask8] = ( hh[mask8] * r[mask8] * f[mask8] - g[mask8] * ap[mask8] + ye[mask8] ) / t[mask8] ye[mask8] = ye[mask8] / t[mask8] - mask8x = box * (p1 >= 0) - if np.any(mask8x): + mask8x = box & (p1 >= 0) + if xp.any(mask8x): de[mask8x] = v[mask8x] / w[mask8x] ye[mask8x] = 0 u[mask8x] = (e[mask8x] + p1[mask8x]) / t[mask8x] v[mask8x] = t[mask8x] / w[mask8x] z[mask8x] = 1.0 - mask9 = box * (s > 1) - if np.any(mask9): + mask9 = box & (s > 1) + if xp.any(mask9): h[mask9] = u[mask9] u[mask9] = v[mask9] v[mask9] = h[mask9] y = 1 / y - e = np.copy(s) - n, t = np.ones((2, len(p))) - m, l = np.zeros((2, len(p))) - - mask10 = np.ones(len(p), dtype=bool) # dynamic mask, changed in each loop iteration - while np.any(mask10): + e = xp.asarray(s, copy=True) + nt = xp.ones((2, (p.shape[0]))) + n, t = (nt[i, ...] for i in range(2)) + ml = xp.zeros((2, (p.shape[0]))) + m, l = (ml[i, ...] for i in range(2)) + + mask10 = xp.ones( + (p.shape[0]), dtype=xp.bool + ) # dynamic mask, changed in each loop iteration + while xp.any(mask10): y[mask10] = y[mask10] - e[mask10] / y[mask10] - mask11 = mask10 * (y == 0.0) - y[mask11] = np.sqrt(e[mask11]) * CB + mask11 = mask10 & (y == 0.0) + y[mask11] = xp.sqrt(e[mask11]) * CB f[mask10] = c[mask10] c[mask10] = d[mask10] / q[mask10] + c[mask10] @@ -454,67 +463,67 @@ def el3v(x, kc, p): t[mask10] = s[mask10] + t[mask10] n[mask10] = 2 * n[mask10] m[mask10] = 2 * m[mask10] - bo10 = mask10 * bo - if np.any(bo10): - bo10b = bo10 * (z < 0) + bo10 = mask10 & bo + if xp.any(bo10): + bo10b = bo10 & (z < 0) m[bo10b] = k[bo10b] + m[bo10b] - k[bo10] = np.sign(r[bo10]) + k[bo10] = xp.sign(r[bo10]) h[bo10] = e[bo10] / (u[bo10] * u[bo10] + v[bo10] * v[bo10]) u[bo10] = u[bo10] * (1 + h[bo10]) v[bo10] = v[bo10] * (1 - h[bo10]) - bo10x = np.array(mask10 * ~bo10, dtype=bool) - if np.any(bo10x): + bo10x = xp.asarray(mask10 & ~bo10, dtype=xp.bool) + if xp.any(bo10x): r[bo10x] = u[bo10x] / v[bo10x] h[bo10x] = z[bo10x] * r[bo10x] z[bo10x] = h[bo10x] * z[bo10x] hh[bo10x] = e[bo10x] / v[bo10x] - bo10x_bk = np.array(bo10x * bk, dtype=bool) # if bk - bo10x_bkx = np.array(bo10x * ~bk, dtype=bool) - if np.any(bo10x_bk): + bo10x_bk = xp.asarray(bo10x * bk, dtype=bool) # if bk + bo10x_bkx = xp.asarray(bo10x * ~bk, dtype=bool) + if xp.any(bo10x_bk): de[bo10x_bk] = de[bo10x_bk] / u[bo10x_bk] ye[bo10x_bk] = ye[bo10x_bk] * (h[bo10x_bk] + 1 / h[bo10x_bk]) + de[ bo10x_bk ] * (1 + r[bo10x_bk]) de[bo10x_bk] = de[bo10x_bk] * (u[bo10x_bk] - hh[bo10x_bk]) - bk[bo10x_bk] = np.abs(ye[bo10x_bk]) < 1 - if np.any(bo10x_bkx): - a_crack = np.log(x[bo10x_bkx]) + bk[bo10x_bk] = xp.abs(ye[bo10x_bk]) < 1 + if xp.any(bo10x_bkx): + a_crack = xp.log(x[bo10x_bkx]) k[bo10x_bkx] = (a_crack / ln2).astype(int) + 1 a_crack = a_crack - k[bo10x_bkx] * ln2 - m[bo10x_bkx] = np.exp(a_crack) + m[bo10x_bkx] = xp.exp(a_crack) m[bo10x_bkx] = m[bo10x_bkx] + k[bo10x_bkx] - mask11 = np.abs(g - s) > CA * g - if np.any(mask11): - bo11 = mask11 * bo - if np.any(bo11): + mask11 = xp.abs(g - s) > CA * g + if xp.any(mask11): + bo11 = mask11 & bo + if xp.any(bo11): g[bo11] = (1 / r[bo11] - r[bo11]) * 0.5 hh[bo11] = u[bo11] + v[bo11] * g[bo11] h[bo11] = g[bo11] * u[bo11] - v[bo11] - bo11b = bo11 * (hh == 0) + bo11b = bo11 & (hh == 0) hh[bo11b] = u[bo11b] * CB - bo11c = bo11 * (h == 0) + bo11c = bo11 & (h == 0) h[bo11c] = v[bo11c] * CB z[bo11] = r[bo11] * h[bo11] r[bo11] = hh[bo11] / h[bo11] - bo11x = mask11 * ~bo - if np.any(bo11x): + bo11x = mask11 & ~bo + if xp.any(bo11x): u[bo11x] = u[bo11x] + e[bo11x] / u[bo11x] v[bo11x] = v[bo11x] + hh[bo11x] - s[mask11] = np.sqrt(e[mask11]) + s[mask11] = xp.sqrt(e[mask11]) s[mask11] = 2 * s[mask11] e[mask11] = s[mask11] * t[mask11] l[mask11] = 2 * l[mask11] - mask12 = mask11 * (y < 0) + mask12 = mask11 & (y < 0) l[mask12] = l[mask12] + 1 # break off parts that have completed their iteration @@ -523,33 +532,33 @@ def el3v(x, kc, p): mask12 = y < 0 l[mask12] = l[mask12] + 1 - e = np.arctan(t / y) + np.pi * l + e = xp.atan(t / y) + xp.pi * l e = e * (c * t + d) / (t * (t + q)) - if np.any(bo): + if xp.any(bo): h[bo] = v[bo] / (t[bo] + u[bo]) z[bo] = 1 - r[bo] * h[bo] h[bo] = r[bo] + h[bo] - bob = bo * (z == 0) + bob = bo & (z == 0) z[bob] = CB - boc = bo * (z < 0) - m[boc] = m[boc] + np.sign(h[boc]) + boc = bo & (z < 0) + m[boc] = m[boc] + xp.sign(h[boc]) - s[bo] = np.arctan(h[bo] / z[bo]) + m[bo] * np.pi + s[bo] = xp.atan(h[bo] / z[bo]) + m[bo] * xp.pi box = ~bo - if np.any(box): - box_bk = box * bk - s[box_bk] = np.arcsinh(ye[box_bk]) + if xp.any(box): + box_bk = box & bk + s[box_bk] = xp.asinh(ye[box_bk]) - box_bkx = box * ~bk - s[box_bkx] = np.log(z[box_bkx]) + m[box_bkx] * ln2 + box_bkx = box & ~bk + s[box_bkx] = xp.log(z[box_bkx]) + m[box_bkx] * ln2 s[box] = s[box] * 0.5 - e = (e + np.sqrt(fa) * s) / n - result[~mask2] = np.sign(x) * e + e = (e + xp.sqrt(fa) * s) / n + result[~mask2] = xp.sign(x) * e # include mask0-case result0[mask0] = result @@ -561,10 +570,11 @@ def el3(xv: np.ndarray, kcv: np.ndarray, pv: np.ndarray) -> np.ndarray: """ combine vectorized and non-vectorized implementations for improved performance """ - n_input = len(xv) + xp = array_namespace(xv, kcv, pv) + n_input = xv.shape[0] - if n_input < 10: - return np.array([el30(x, kc, p) for x, kc, p in zip(xv, kcv, pv, strict=False)]) + # if n_input < 10: + # return xp.asarray([el30(x, kc, p) for x, kc, p in zip(xv, kcv, pv, strict=False)]) return el3v(xv, kcv, pv) @@ -584,72 +594,73 @@ def el3_angle(phi: np.ndarray, n: np.ndarray, m: np.ndarray) -> np.ndarray: -> possibly provide a non-vectorized form of this ? (maybe not worth the effort) """ # pylint: disable=too-many-statements - n_vec = len(phi) - results = np.zeros(n_vec) + xp = array_namespace(phi, n, m) + n_vec = phi.shape[0] + results = xp.zeros(n_vec) - kc = np.sqrt(1 - m) + kc = xp.sqrt(1 - m) p = 1 - n D = 8 - n = (phi / np.pi).astype(int) - phi_red = phi - n * np.pi + n = xp.round(phi / xp.pi) + phi_red = phi - n * xp.pi - mask1 = (n <= 0) * (phi_red < -np.pi / 2) - mask2 = (n >= 0) * (phi_red > np.pi / 2) - if np.any(mask1): + mask1 = (n <= 0) & (phi_red < -xp.pi / 2) + mask2 = (n >= 0) & (phi_red > xp.pi / 2) + if xp.any(mask1): n[mask1] = n[mask1] - 1 - phi_red[mask1] = phi_red[mask1] + np.pi + phi_red[mask1] = phi_red[mask1] + xp.pi - if np.any(mask2): + if xp.any(mask2): n[mask2] = n[mask2] + 1 - phi_red[mask2] = phi_red[mask2] - np.pi + phi_red[mask2] = phi_red[mask2] - xp.pi mask3 = n != 0 mask3x = ~mask3 - if np.any(mask3): + if xp.any(mask3): n3, phi3, p3, kc3 = n[mask3], phi[mask3], p[mask3], kc[mask3] phi_red3 = phi_red[mask3] - results3 = np.zeros(np.sum(mask3)) - onez = np.ones(np.sum(mask3)) + results3 = xp.zeros(xp.count_nonzero(mask3)) + onez = xp.ones(xp.count_nonzero(mask3)) cel3_res = cel(kc3, p3, onez, onez) # 3rd kind cel - mask3a = phi_red3 > np.pi / 2 - 10 ** (-D) - mask3b = phi_red3 < -np.pi / 2 + 10 ** (-D) - mask3c = ~mask3a * ~mask3b + mask3a = phi_red3 > xp.pi / 2 - 10 ** (-D) + mask3b = phi_red3 < -xp.pi / 2 + 10 ** (-D) + mask3c = ~mask3a & ~mask3b - if np.any(mask3a): + if xp.any(mask3a): results3[mask3a] = (2 * n3[mask3a] + 1) * cel3_res[mask3a] - if np.any(mask3b): + if xp.any(mask3b): results3[mask3b] = (2 * n3[mask3b] - 1) * cel3_res[mask3b] - if np.any(mask3c): - el33_res = el3(np.tan(phi3[mask3c]), kc3[mask3c], p3[mask3c]) + if xp.any(mask3c): + el33_res = el3(xp.tan(phi3[mask3c]), kc3[mask3c], p3[mask3c]) results3[mask3c] = 2 * n3[mask3c] * cel3_res[mask3c] + el33_res results[mask3] = results3 - if np.any(mask3x): + if xp.any(mask3x): phi_red3x = phi_red[mask3x] - results3x = np.zeros(np.sum(mask3x)) + results3x = xp.zeros(xp.count_nonzero(mask3x)) phi3x, kc3x, p3x = phi[mask3x], kc[mask3x], p[mask3x] - mask3xa = phi_red3x > np.pi / 2 - 10 ** (-D) - mask3xb = phi_red3x < -np.pi / 2 + 10 ** (-D) - mask3xc = ~mask3xa * ~mask3xb + mask3xa = phi_red3x > xp.pi / 2 - 10 ** (-D) + mask3xb = phi_red3x < -xp.pi / 2 + 10 ** (-D) + mask3xc = ~mask3xa & ~mask3xb - if np.any(mask3xa): - onez = np.ones(np.sum(mask3xa)) + if xp.any(mask3xa): + onez = xp.ones(xp.count_nonzero(mask3xa)) results3x[mask3xa] = cel( kc3x[mask3xa], p3x[mask3xa], onez, onez ) # 3rd kind cel - if np.any(mask3xb): - onez = np.ones(np.sum(mask3xb)) + if xp.any(mask3xb): + onez = xp.ones(xp.count_nonzero(mask3xb)) results3x[mask3xb] = -cel( kc3x[mask3xb], p3x[mask3xb], onez, onez ) # 3rd kind cel - if np.any(mask3xc): + if xp.any(mask3xc): results3x[mask3xc] = el3( - np.tan(phi3x[mask3xc]), kc3x[mask3xc], p3x[mask3xc] + xp.tan(phi3x[mask3xc]), kc3x[mask3xc], p3x[mask3xc] ) results[mask3x] = results3x diff --git a/src/magpylib/_src/fields/special_elliptic.py b/src/magpylib/_src/fields/special_elliptic.py new file mode 100644 index 000000000..6aad39794 --- /dev/null +++ b/src/magpylib/_src/fields/special_elliptic.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from typing import Any + +import array_api_extra as xpx +import scipy as sp + +type Array = Any + + +def ellipeinc(phi: Array, m: Array): + return xpx.lazy_apply(sp.special.ellipeinc, phi, m, as_numpy=True) + + +def ellipkinc(phi: Array, m: Array): + return xpx.lazy_apply(sp.special.ellipkinc, phi, m, as_numpy=True) + + +def ellipe(m: Array): + return xpx.lazy_apply(sp.special.ellipe, m, as_numpy=True) + + +def ellipk(m: Array): + return xpx.lazy_apply(sp.special.ellipk, m, as_numpy=True) From 97368e2975e1dd8984d18250c6d7f3d52ba9bb24 Mon Sep 17 00:00:00 2001 From: purepani Date: Mon, 19 May 2025 14:00:16 -0500 Subject: [PATCH 13/16] Adds array api support for cylinder core functions --- src/magpylib/_src/fields/field_BH_cylinder.py | 81 +++++++++++-------- 1 file changed, 47 insertions(+), 34 deletions(-) diff --git a/src/magpylib/_src/fields/field_BH_cylinder.py b/src/magpylib/_src/fields/field_BH_cylinder.py index 04fa1a52a..8a8457edc 100644 --- a/src/magpylib/_src/fields/field_BH_cylinder.py +++ b/src/magpylib/_src/fields/field_BH_cylinder.py @@ -6,13 +6,18 @@ # pylint: disable = no-name-in-module from __future__ import annotations +import array_api_extra as xpx import numpy as np + +# from scipy.special import ellipe, ellipk +from array_api_compat import array_namespace from scipy.constants import mu_0 as MU0 -from scipy.special import ellipe, ellipk from magpylib._src.fields.special_cel import cel +from magpylib._src.fields.special_elliptic import ellipe, ellipk from magpylib._src.input_checks import check_field_input from magpylib._src.utility import cart_to_cyl_coordinates, cyl_field_to_cart +from magpylib._src.array_api_utils import xp_promote # CORE @@ -58,22 +63,27 @@ def magnet_cylinder_axial_Bfield(z0: np.ndarray, r: np.ndarray, z: np.ndarray) - ----- Implementation based on Derby, American Journal of Physics 78.3 (2010): 229-235. """ - n = len(z0) + xp = array_namespace(z0, r, z) + z0, r, z = xp_promote(z0, r, z, force_floating=True, xp=xp) + n = z0.shape[0] + z0 = xp.astype(z0, xp.float64) + r = xp.astype(r, xp.float64) + z = xp.astype(z, xp.float64) # some important quantities zph, zmh = z + z0, z - z0 dpr, dmr = 1 + r, 1 - r - sq0 = np.sqrt(zmh**2 + dpr**2) - sq1 = np.sqrt(zph**2 + dpr**2) + sq0 = xp.sqrt(zmh**2 + dpr**2) + sq1 = xp.sqrt(zph**2 + dpr**2) - k1 = np.sqrt((zph**2 + dmr**2) / (zph**2 + dpr**2)) - k0 = np.sqrt((zmh**2 + dmr**2) / (zmh**2 + dpr**2)) + k1 = xp.sqrt((zph**2 + dmr**2) / (zph**2 + dpr**2)) + k0 = xp.sqrt((zmh**2 + dmr**2) / (zmh**2 + dpr**2)) gamma = dmr / dpr - one = np.ones(n) + one = xp.ones(n, dtype=xp.float64) # radial field (unit polarization) - Br = (cel(k1, one, one, -one) / sq1 - cel(k0, one, one, -one) / sq0) / np.pi + Br = (cel(k1, one, one, -one) / sq1 - cel(k0, one, one, -one) / sq0) / xp.pi # axial field (unit polarization) Bz = ( @@ -83,10 +93,10 @@ def magnet_cylinder_axial_Bfield(z0: np.ndarray, r: np.ndarray, z: np.ndarray) - zph * cel(k1, gamma**2, one, gamma) / sq1 - zmh * cel(k0, gamma**2, one, gamma) / sq0 ) - / np.pi + / xp.pi ) - return np.vstack((Br, np.zeros(n), Bz)) + return xp.stack((Br, xp.zeros(n), Bz)) # CORE @@ -145,11 +155,14 @@ def magnet_cylinder_diametral_Hfield( (unpublished). """ # pylint: disable=too-many-statements + xp = array_namespace(z0, r, z, phi) + z0, r, z, phi = xp_promote(z0, r, z, phi, force_floating=True, xp=xp) - n = len(z0) + n = z0.shape[0] # allocate to treat small r special cases - Hr, Hphi, Hz = np.empty((3, n)) + H = xp.empty((3, n)) + Hr, Hphi, Hz = H[0, ...], H[1, ...], H[2, ...] # compute repeated quantities for all cases zp = z + z0 @@ -162,7 +175,7 @@ def magnet_cylinder_diametral_Hfield( # case small_r: numerical instability of general solution mask_small_r = r < 0.05 mask_general = ~mask_small_r - if np.any(mask_small_r): + if xp.any(mask_small_r): phiX = phi[mask_small_r] zpX, zmX = zp[mask_small_r], zm[mask_small_r] zp2X, zm2X = zp2[mask_small_r], zm2[mask_small_r] @@ -171,8 +184,8 @@ def magnet_cylinder_diametral_Hfield( # taylor series for small r zpp = zp2X + 1 zmm = zm2X + 1 - sqrt_p = np.sqrt(zpp) - sqrt_m = np.sqrt(zmm) + sqrt_p = xp.sqrt(zpp) + sqrt_m = xp.sqrt(zmm) frac1 = zpX / sqrt_p frac2 = zmX / sqrt_m @@ -189,12 +202,12 @@ def magnet_cylinder_diametral_Hfield( * r4X ) - Hr[mask_small_r] = -np.cos(phiX) / 4 * (term1 + 9 * term2 + 25 * term3) + Hr[mask_small_r] = -xp.cos(phiX) / 4 * (term1 + 9 * term2 + 25 * term3) - Hphi[mask_small_r] = np.sin(phiX) / 4 * (term1 + 3 * term2 + 5 * term3) + Hphi[mask_small_r] = xp.sin(phiX) / 4 * (term1 + 3 * term2 + 5 * term3) Hz[mask_small_r] = ( - -np.cos(phiX) + -xp.cos(phiX) / 4 * ( rX * (1 / zpp / sqrt_p - 1 / zmm / sqrt_m) @@ -215,12 +228,12 @@ def magnet_cylinder_diametral_Hfield( # if there are small_r, select the general/case variables # when there are no small_r cases it is not necessary to slice with [True, True, Tue,...] phi = phi[mask_general] - n = len(phi) + n = phi.shape[0] zp, zm = zp[mask_general], zm[mask_general] zp2, zm2 = zp2[mask_general], zm2[mask_general] r, r2 = r[mask_general], r2[mask_general] - if np.any(mask_general): + if xp.any(mask_general): rp = r + 1 rm = r - 1 rp2 = rp**2 @@ -228,8 +241,8 @@ def magnet_cylinder_diametral_Hfield( ap2 = zp2 + rm**2 am2 = zm2 + rm**2 - ap = np.sqrt(ap2) - am = np.sqrt(am2) + ap = xp.sqrt(ap2) + am = xp.sqrt(am2) argp = -4 * r / ap2 argm = -4 * r / am2 @@ -238,24 +251,24 @@ def magnet_cylinder_diametral_Hfield( # result is numerically stable in the vicinity of of r=r0 # so only the special case must be caught (not the surroundings) mask_special = rm == 0 - argc = np.ones(n) * 1e16 # should be np.Inf but leads to 1/0 problems in cel + argc = xp.ones(n) * 1e16 # should be np.Inf but leads to 1/0 problems in cel argc[~mask_special] = -4 * r[~mask_special] / rm2[~mask_special] # special case 1/rm - one_over_rm = np.zeros(n) + one_over_rm = xp.zeros(n) one_over_rm[~mask_special] = 1 / rm[~mask_special] elle_p = ellipe(argp) elle_m = ellipe(argm) ellk_p = ellipk(argp) ellk_m = ellipk(argm) - onez = np.ones(n) - ellpi_p = cel(np.sqrt(1 - argp), 1 - argc, onez, onez) # elliptic_Pi - ellpi_m = cel(np.sqrt(1 - argm), 1 - argc, onez, onez) # elliptic_Pi + onez = xp.ones(n) + ellpi_p = cel(xp.sqrt(1 - argp), 1 - argc, onez, onez) # elliptic_Pi + ellpi_m = cel(xp.sqrt(1 - argm), 1 - argc, onez, onez) # elliptic_Pi # compute fields Hr[mask_general] = ( - -np.cos(phi) - / (4 * np.pi * r2) + -xp.cos(phi) + / (4 * xp.pi * r2) * ( -zm * am * elle_m + zp * ap * elle_p @@ -266,8 +279,8 @@ def magnet_cylinder_diametral_Hfield( ) Hphi[mask_general] = ( - np.sin(phi) - / (4 * np.pi * r2) + xp.sin(phi) + / (4 * xp.pi * r2) * ( +zm * am * elle_m - zp * ap * elle_p @@ -279,8 +292,8 @@ def magnet_cylinder_diametral_Hfield( ) Hz[mask_general] = ( - -np.cos(phi) - / (2 * np.pi * r) + -xp.cos(phi) + / (2 * xp.pi * r) * ( +am * elle_m - ap * elle_p @@ -289,7 +302,7 @@ def magnet_cylinder_diametral_Hfield( ) ) - return np.vstack((Hr, Hphi, Hz)) + return xp.stack((Hr, Hphi, Hz)) def BHJM_magnet_cylinder( From 0fac948a7e34556e8d873fd78ab12d96faad97df Mon Sep 17 00:00:00 2001 From: purepani Date: Mon, 19 May 2025 15:59:27 -0500 Subject: [PATCH 14/16] Adds array api support to cylinder segment core function --- .../_src/fields/field_BH_cylinder_segment.py | 1801 +++++++++-------- 1 file changed, 911 insertions(+), 890 deletions(-) diff --git a/src/magpylib/_src/fields/field_BH_cylinder_segment.py b/src/magpylib/_src/fields/field_BH_cylinder_segment.py index 2d0e775b7..bfb03724e 100644 --- a/src/magpylib/_src/fields/field_BH_cylinder_segment.py +++ b/src/magpylib/_src/fields/field_BH_cylinder_segment.py @@ -6,13 +6,16 @@ # pylint: disable=too-many-positional-arguments from __future__ import annotations +import array_api_extra as xpx import numpy as np +from array_api_compat import array_namespace from scipy.constants import mu_0 as MU0 -from scipy.special import ellipeinc, ellipkinc from magpylib._src.fields.field_BH_cylinder import BHJM_magnet_cylinder from magpylib._src.fields.special_el3 import el3_angle +from magpylib._src.fields.special_elliptic import ellipeinc, ellipkinc from magpylib._src.input_checks import check_field_input +from magpylib._src.array_api_utils import xp_promote def arctan_k_tan_2(k, phi): @@ -23,16 +26,18 @@ def arctan_k_tan_2(k, phi): can be replaced by non-masked version ? """ + xp = array_namespace(k, phi) - full_periods = np.round(phi / (2.0 * np.pi)) - phi_red = phi - full_periods * 2.0 * np.pi + full_periods = xp.round(phi / (2.0 * xp.pi)) + phi_red = phi - full_periods * 2.0 * xp.pi - result = full_periods * np.pi + result = full_periods * xp.pi - return np.where( - np.abs(phi_red) < np.pi, - result + np.arctan(k * np.tan(phi_red / 2.0)), - result + phi_red / 2.0, + return xpx.apply_where( + xp.abs(phi_red) < xp.pi, + (result, k, phi_red), + lambda result, k, phi_red: result + xp.atan(k * xp.tan(phi_red / 2.0)), + lambda result, k, phi_red: result + phi_red / 2.0, ) @@ -42,7 +47,7 @@ def close(arg1: np.ndarray, arg2: np.ndarray) -> np.ndarray: input: ndarray, shape (n,) or numpy-interpretable scalar output: ndarray, dtype=bool """ - return np.isclose(arg1, arg2, rtol=1e-12, atol=1e-12) + return xpx.isclose(arg1, arg2, rtol=1e-12, atol=1e-12) def determine_cases(r, phi, z, r1, phi1, z1): @@ -60,37 +65,38 @@ def determine_cases(r, phi, z, r1, phi1, z1): 2nd digit: 1:phi-phi1= 2n*pi, 2:phi-phi1=(2n+1)*pi, 3:general 3rd digit: 1:r=r1=0, 2:r=0, 3:r1=0, 4:r=r1>0, 5:general """ - n = len(r) # input length + xp = array_namespace(r, phi, z, r1, phi1, z1) + n = r.shape[0] # input length # allocate result - result = np.ones((3, n)) + result = xp.ones((3, n)) # identify z-case mask_z = close(z, z1) - result[0] = 200 - result[0, mask_z] = 100 + result[0, ...] = 200 + result[0, :][mask_z] = 100 # identify phi-case - mod_2pi = np.abs(phi - phi1) % (2 * np.pi) - mask_phi1 = np.logical_or(close(mod_2pi, 0), close(mod_2pi, 2 * np.pi)) - mod_pi = np.abs(phi - phi1) % np.pi - mask_phi2 = np.logical_or(close(mod_pi, 0), close(mod_pi, np.pi)) - result[1] = 30 - result[1, mask_phi2] = 20 - result[1, mask_phi1] = 10 + mod_2pi = xp.abs(phi - phi1) % (2 * xp.pi) + mask_phi1 = xp.logical_or(close(mod_2pi, 0), close(mod_2pi, 2 * xp.pi)) + mod_pi = xp.abs(phi - phi1) % xp.pi + mask_phi2 = xp.logical_or(close(mod_pi, 0), close(mod_pi, xp.pi)) + result[1, ...] = 30 + result[1, ...][mask_phi2] = 20 + result[1, ...][mask_phi1] = 10 # identify r-case mask_r2 = close(r, 0) mask_r3 = close(r1, 0) mask_r4 = close(r, r1) - mask_r1 = mask_r2 * mask_r3 - result[2] = 5 - result[2, mask_r4] = 4 - result[2, mask_r3] = 3 - result[2, mask_r2] = 2 - result[2, mask_r1] = 1 + mask_r1 = mask_r2 & mask_r3 + result[2, ...] = 5 + result[2, ...][mask_r4] = 4 + result[2, ...][mask_r3] = 3 + result[2, ...][mask_r2] = 2 + result[2, ...][mask_r1] = 1 - return np.array(np.sum(result, axis=0), dtype=int) + return xp.asarray(xp.sum(result, axis=0), dtype=xp.int32) # Implementation of all non-zero field components in every special case @@ -100,384 +106,382 @@ def determine_cases(r, phi, z, r1, phi1, z1): # 112 ############## -def Hphi_zk_case112(r_i, theta_M): - return np.cos(theta_M) * np.log(r_i) +def Hphi_zk_case112(xp, r_i, theta_M): + return xp.cos(theta_M) * xp.log(r_i) -def Hz_ri_case112(phi_bar_M, theta_M): - return -np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_ri_case112(xp, phi_bar_M, theta_M): + return -xp.sin(theta_M) * xp.sin(phi_bar_M) -def Hz_phij_case112(r_i, phi_bar_M, theta_M): - return np.sin(theta_M) * np.sin(phi_bar_M) * np.log(r_i) +def Hz_phij_case112(xp, r_i, phi_bar_M, theta_M): + return xp.sin(theta_M) * xp.sin(phi_bar_M) * xp.log(r_i) # 113 ############## -def Hphi_zk_case113(r, theta_M): - return -np.cos(theta_M) * np.log(r) +def Hphi_zk_case113(xp, r, theta_M): + return -xp.cos(theta_M) * xp.log(r) -def Hz_phij_case113(r, phi_bar_M, theta_M): - return -np.sin(theta_M) * np.sin(phi_bar_M) * np.log(r) +def Hz_phij_case113(xp, r, phi_bar_M, theta_M): + return -xp.sin(theta_M) * xp.sin(phi_bar_M) * xp.log(r) # 115 ############## -def Hr_zk_case115(r, r_i, r_bar_i, phi_bar_j, theta_M): +def Hr_zk_case115(xp, r, r_i, r_bar_i, phi_bar_j, theta_M): E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) - E_coef = np.cos(theta_M) * np.abs(r_bar_i) / r + E_coef = xp.cos(theta_M) * xp.abs(r_bar_i) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) - F_coef = -np.cos(theta_M) * (r**2 + r_i**2) / (r * np.abs(r_bar_i)) + F_coef = -xp.cos(theta_M) * (r**2 + r_i**2) / (r * xp.abs(r_bar_i)) return E_coef * E + F_coef * F -def Hphi_zk_case115(r, r_i, r_bar_i, theta_M): +def Hphi_zk_case115(xp, r, r_i, r_bar_i, theta_M): t1 = r_i / r - t1_coef = -np.cos(theta_M) * np.sign(r_bar_i) - t2 = np.log(np.abs(r_bar_i)) * np.sign(r_bar_i) - t2_coef = -np.cos(theta_M) + t1_coef = -xp.cos(theta_M) * xp.sign(r_bar_i) + t2 = xp.log(xp.abs(r_bar_i)) * xp.sign(r_bar_i) + t2_coef = -xp.cos(theta_M) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case115(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): - t1 = np.abs(r_bar_i) / r - t1_coef = np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_ri_case115(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): + t1 = xp.abs(r_bar_i) / r + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) - E_coef = np.sin(theta_M) * np.cos(phi_bar_M) * np.abs(r_bar_i) / r + E_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.abs(r_bar_i) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) F_coef = ( - -np.sin(theta_M) * np.cos(phi_bar_M) * (r**2 + r_i**2) / (r * np.abs(r_bar_i)) + -xp.sin(theta_M) * xp.cos(phi_bar_M) * (r**2 + r_i**2) / (r * xp.abs(r_bar_i)) ) return t1_coef * t1 + E_coef * E + F_coef * F -def Hz_phij_case115(r_bar_i, phi_bar_M, theta_M): - t1 = np.log(np.abs(r_bar_i)) * np.sign(r_bar_i) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_phij_case115(xp, r_bar_i, phi_bar_M, theta_M): + t1 = xp.log(xp.abs(r_bar_i)) * xp.sign(r_bar_i) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 # 122 ############## -def Hphi_zk_case122(r_i, theta_M): - return -np.cos(theta_M) * np.log(r_i) +def Hphi_zk_case122(xp, r_i, theta_M): + return -xp.cos(theta_M) * xp.log(r_i) -def Hz_ri_case122(phi_bar_M, theta_M): - return np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_ri_case122(xp, phi_bar_M, theta_M): + return xp.sin(theta_M) * xp.sin(phi_bar_M) -def Hz_phij_case122(r_i, phi_bar_M, theta_M): - return -np.sin(theta_M) * np.sin(phi_bar_M) * np.log(r_i) +def Hz_phij_case122(xp, r_i, phi_bar_M, theta_M): + return -xp.sin(theta_M) * xp.sin(phi_bar_M) * xp.log(r_i) # 123 ############## -def Hphi_zk_case123(r, theta_M): - return -np.cos(theta_M) * np.log(r) +def Hphi_zk_case123(xp, r, theta_M): + return -xp.cos(theta_M) * xp.log(r) -def Hz_phij_case123(r, phi_bar_M, theta_M): - return -np.sin(theta_M) * np.sin(phi_bar_M) * np.log(r) +def Hz_phij_case123(xp, r, phi_bar_M, theta_M): + return -xp.sin(theta_M) * xp.sin(phi_bar_M) * xp.log(r) # 124 ############## -def Hphi_zk_case124(r, theta_M): - return np.cos(theta_M) * (1.0 - np.log(2.0 * r)) +def Hphi_zk_case124(xp, r, theta_M): + return xp.cos(theta_M) * (1.0 - xp.log(2.0 * r)) -def Hz_ri_case124(phi_bar_M, theta_M): - return 2.0 * np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_ri_case124(xp, phi_bar_M, theta_M): + return 2.0 * xp.sin(theta_M) * xp.sin(phi_bar_M) -def Hz_phij_case124(r, phi_bar_M, theta_M): - return -np.sin(theta_M) * np.sin(phi_bar_M) * np.log(2.0 * r) +def Hz_phij_case124(xp, r, phi_bar_M, theta_M): + return -xp.sin(theta_M) * xp.sin(phi_bar_M) * xp.log(2.0 * r) # 125 ############## -def Hr_zk_case125(r, r_i, r_bar_i, phi_bar_j, theta_M): +def Hr_zk_case125(xp, r, r_i, r_bar_i, phi_bar_j, theta_M): E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) - E_coef = np.cos(theta_M) * np.abs(r_bar_i) / r + E_coef = xp.cos(theta_M) * xp.abs(r_bar_i) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) - F_coef = -np.cos(theta_M) * (r**2 + r_i**2) / (r * np.abs(r_bar_i)) + F_coef = -xp.cos(theta_M) * (r**2 + r_i**2) / (r * xp.abs(r_bar_i)) return E_coef * E + F_coef * F -def Hphi_zk_case125(r, r_i, theta_M): - return np.cos(theta_M) / r * (r_i - r * np.log(r + r_i)) +def Hphi_zk_case125(xp, r, r_i, theta_M): + return xp.cos(theta_M) / r * (r_i - r * xp.log(r + r_i)) -def Hz_ri_case125(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): +def Hz_ri_case125(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) - E_coef = np.sin(theta_M) * np.cos(phi_bar_M) * np.abs(r_bar_i) / r + E_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.abs(r_bar_i) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) F_coef = ( - -np.sin(theta_M) * np.cos(phi_bar_M) * (r**2 + r_i**2) / (r * np.abs(r_bar_i)) + -xp.sin(theta_M) * xp.cos(phi_bar_M) * (r**2 + r_i**2) / (r * xp.abs(r_bar_i)) ) - return np.sin(theta_M) * np.sin(phi_bar_M) * (r + r_i) / r + E_coef * E + F_coef * F + return xp.sin(theta_M) * xp.sin(phi_bar_M) * (r + r_i) / r + E_coef * E + F_coef * F -def Hz_phij_case125(r, r_i, phi_bar_M, theta_M): - return -np.sin(theta_M) * np.sin(phi_bar_M) * np.log(r + r_i) +def Hz_phij_case125(xp, r, r_i, phi_bar_M, theta_M): + return -xp.sin(theta_M) * xp.sin(phi_bar_M) * xp.log(r + r_i) # 132 ############## -def Hr_zk_case132(r_i, phi_bar_j, theta_M): - return np.cos(theta_M) * np.sin(phi_bar_j) * np.log(r_i) +def Hr_zk_case132(xp, r_i, phi_bar_j, theta_M): + return xp.cos(theta_M) * xp.sin(phi_bar_j) * xp.log(r_i) -def Hphi_zk_case132(r_i, phi_bar_j, theta_M): - return np.cos(theta_M) * np.cos(phi_bar_j) * np.log(r_i) +def Hphi_zk_case132(xp, r_i, phi_bar_j, theta_M): + return xp.cos(theta_M) * xp.cos(phi_bar_j) * xp.log(r_i) -def Hz_ri_case132(phi_bar_Mj, theta_M): - return -np.sin(theta_M) * np.sin(phi_bar_Mj) +def Hz_ri_case132(xp, phi_bar_Mj, theta_M): + return -xp.sin(theta_M) * xp.sin(phi_bar_Mj) -def Hz_phij_case132(r_i, phi_bar_Mj, theta_M): - return np.sin(theta_M) * np.sin(phi_bar_Mj) * np.log(r_i) +def Hz_phij_case132(xp, r_i, phi_bar_Mj, theta_M): + return xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.log(r_i) # 133 ############## -def Hr_zk_case133(r, phi_bar_j, theta_M): - return -np.cos(theta_M) * np.sin(phi_bar_j) + np.cos(theta_M) * np.sin( +def Hr_zk_case133(xp, r, phi_bar_j, theta_M): + return -xp.cos(theta_M) * xp.sin(phi_bar_j) + xp.cos(theta_M) * xp.sin( phi_bar_j - ) * np.log(r * (1.0 - np.cos(phi_bar_j))) + ) * xp.log(r * (1.0 - xp.cos(phi_bar_j))) -def Hphi_zk_case133(phi_bar_j, theta_M): - return np.cos(theta_M) - np.cos(theta_M) * np.cos(phi_bar_j) * np.arctanh( - np.cos(phi_bar_j) +def Hphi_zk_case133(xp, phi_bar_j, theta_M): + return xp.cos(theta_M) - xp.cos(theta_M) * xp.cos(phi_bar_j) * xp.atanh( + xp.cos(phi_bar_j) ) -def Hz_phij_case133(phi_bar_j, phi_bar_Mj, theta_M): - return -np.sin(theta_M) * np.sin(phi_bar_Mj) * np.arctanh(np.cos(phi_bar_j)) +def Hz_phij_case133(xp, phi_bar_j, phi_bar_Mj, theta_M): + return -xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.atanh(xp.cos(phi_bar_j)) # 134 ############## -def Hr_zk_case134(r, phi_bar_j, theta_M): - t1 = np.sin(phi_bar_j) - t1_coef = -np.cos(theta_M) - t2 = np.sin(phi_bar_j) / np.sqrt(1.0 - np.cos(phi_bar_j)) - t2_coef = -np.sqrt(2.0) * np.cos(theta_M) - t3 = np.log( - r * (1.0 - np.cos(phi_bar_j) + np.sqrt(2.0) * np.sqrt(1.0 - np.cos(phi_bar_j))) +def Hr_zk_case134(xp, r, phi_bar_j, theta_M): + t1 = xp.sin(phi_bar_j) + t1_coef = -xp.cos(theta_M) + t2 = xp.sin(phi_bar_j) / xp.sqrt(1.0 - xp.cos(phi_bar_j)) + t2_coef = -xp.sqrt(2.0) * xp.cos(theta_M) + t3 = xp.log( + r * (1.0 - xp.cos(phi_bar_j) + xp.sqrt(2.0) * xp.sqrt(1.0 - xp.cos(phi_bar_j))) ) - t3_coef = np.cos(theta_M) * np.sin(phi_bar_j) - t4 = np.arctanh( - np.sin(phi_bar_j) / (np.sqrt(2.0) * np.sqrt(1.0 - np.cos(phi_bar_j))) - ) - t4_coef = np.cos(theta_M) + t3_coef = xp.cos(theta_M) * xp.sin(phi_bar_j) + t4 = xp.atanh(xp.sin(phi_bar_j) / (xp.sqrt(2.0) * xp.sqrt(1.0 - xp.cos(phi_bar_j)))) + t4_coef = xp.cos(theta_M) return t1_coef * t1 + t2_coef * t2 + t3_coef * t3 + t4_coef * t4 -def Hphi_zk_case134(phi_bar_j, theta_M): - return np.sqrt(2) * np.cos(theta_M) * np.sqrt(1 - np.cos(phi_bar_j)) + np.cos( +def Hphi_zk_case134(xp, phi_bar_j, theta_M): + return xp.sqrt(2) * xp.cos(theta_M) * xp.sqrt(1 - xp.cos(phi_bar_j)) + xp.cos( theta_M - ) * np.cos(phi_bar_j) * np.arctanh(np.sqrt((1 - np.cos(phi_bar_j)) / 2)) + ) * xp.cos(phi_bar_j) * xp.atanh(xp.sqrt((1 - xp.cos(phi_bar_j)) / 2)) -def Hz_ri_case134(phi_bar_j, phi_bar_M, theta_M): - t1 = np.sqrt(1.0 - np.cos(phi_bar_j)) - t1_coef = np.sqrt(2.0) * np.sin(theta_M) * np.sin(phi_bar_M) - t2 = np.sin(phi_bar_j) / t1 - t2_coef = -np.sqrt(2.0) * np.sin(theta_M) * np.cos(phi_bar_M) - t3 = np.arctanh(t2 / np.sqrt(2.0)) - t3_coef = np.sin(theta_M) * np.cos(phi_bar_M) +def Hz_ri_case134(xp, phi_bar_j, phi_bar_M, theta_M): + t1 = xp.sqrt(1.0 - xp.cos(phi_bar_j)) + t1_coef = xp.sqrt(2.0) * xp.sin(theta_M) * xp.sin(phi_bar_M) + t2 = xp.sin(phi_bar_j) / t1 + t2_coef = -xp.sqrt(2.0) * xp.sin(theta_M) * xp.cos(phi_bar_M) + t3 = xp.atanh(t2 / xp.sqrt(2.0)) + t3_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) return t1_coef * t1 + t2_coef * t2 + t3_coef * t3 -def Hz_phij_case134(phi_bar_j, phi_bar_Mj, theta_M): +def Hz_phij_case134(xp, phi_bar_j, phi_bar_Mj, theta_M): return ( - np.sin(theta_M) - * np.sin(phi_bar_Mj) - * np.arctanh(np.sqrt((1.0 - np.cos(phi_bar_j)) / 2.0)) + xp.sin(theta_M) + * xp.sin(phi_bar_Mj) + * xp.atanh(xp.sqrt((1.0 - xp.cos(phi_bar_j)) / 2.0)) ) # 135 ############## -def Hr_zk_case135(r, r_i, r_bar_i, phi_bar_j, theta_M): - t1 = np.sin(phi_bar_j) - t1_coef = -np.cos(theta_M) - t2 = np.log( +def Hr_zk_case135(xp, r, r_i, r_bar_i, phi_bar_j, theta_M): + t1 = xp.sin(phi_bar_j) + t1_coef = -xp.cos(theta_M) + t2 = xp.log( r_i - - r * np.cos(phi_bar_j) - + np.sqrt(r_i**2 + r**2 - 2 * r_i * r * np.cos(phi_bar_j)) + - r * xp.cos(phi_bar_j) + + xp.sqrt(r_i**2 + r**2 - 2 * r_i * r * xp.cos(phi_bar_j)) ) - t2_coef = np.cos(theta_M) * np.sin(phi_bar_j) + t2_coef = xp.cos(theta_M) * xp.sin(phi_bar_j) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) - E_coef = np.cos(theta_M) * np.abs(r_bar_i) / r + E_coef = xp.cos(theta_M) * xp.abs(r_bar_i) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / r_bar_i**2) - F_coef = -np.cos(theta_M) * (r**2 + r_i**2) / (r * np.abs(r_bar_i)) + F_coef = -xp.cos(theta_M) * (r**2 + r_i**2) / (r * xp.abs(r_bar_i)) return t1_coef * t1 + t2_coef * t2 + E_coef * E + F_coef * F -def Hphi_zk_case135(r, r_i, phi_bar_j, theta_M): - t1 = np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j)) - t1_coef = np.cos(theta_M) / r - t2 = np.arctanh((r * np.cos(phi_bar_j) - r_i) / t1) - t2_coef = -np.cos(theta_M) * np.cos(phi_bar_j) +def Hphi_zk_case135(xp, r, r_i, phi_bar_j, theta_M): + t1 = xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j)) + t1_coef = xp.cos(theta_M) / r + t2 = xp.atanh((r * xp.cos(phi_bar_j) - r_i) / t1) + t2_coef = -xp.cos(theta_M) * xp.cos(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case135(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): +def Hz_ri_case135(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): t = r_bar_i**2 - t1 = np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j)) / r - t1_coef = np.sin(theta_M) * np.sin(phi_bar_M) + t1 = xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j)) / r + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / t) - E_coef = np.sin(theta_M) * np.cos(phi_bar_M) * np.sqrt(t) / r + E_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.sqrt(t) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / t) - F_coef = -np.sin(theta_M) * np.cos(phi_bar_M) * (r**2 + r_i**2) / (r * np.sqrt(t)) + F_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) * (r**2 + r_i**2) / (r * xp.sqrt(t)) return t1_coef * t1 + E_coef * E + F_coef * F -def Hz_phij_case135(r, r_i, phi_bar_j, phi_bar_Mj, theta_M): - t1 = np.arctanh( - (r * np.cos(phi_bar_j) - r_i) - / np.sqrt(r**2 + r_i**2 - 2 * r * r_i * np.cos(phi_bar_j)) +def Hz_phij_case135(xp, r, r_i, phi_bar_j, phi_bar_Mj, theta_M): + t1 = xp.atanh( + (r * xp.cos(phi_bar_j) - r_i) + / xp.sqrt(r**2 + r_i**2 - 2 * r * r_i * xp.cos(phi_bar_j)) ) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) return t1_coef * t1 # 211 ############## -def Hr_phij_case211(phi_bar_M, theta_M, z_bar_k): +def Hr_phij_case211(xp, phi_bar_M, theta_M, z_bar_k): return ( - -np.sin(theta_M) - * np.sin(phi_bar_M) - * np.sign(z_bar_k) - * np.log(np.abs(z_bar_k)) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) + * xp.sign(z_bar_k) + * xp.log(xp.abs(z_bar_k)) ) -def Hz_zk_case211(phi_j, theta_M, z_bar_k): - return -np.cos(theta_M) * np.sign(z_bar_k) * phi_j +def Hz_zk_case211(xp, phi_j, theta_M, z_bar_k): + return -xp.cos(theta_M) * xp.sign(z_bar_k) * phi_j # 212 ############## -def Hr_ri_case212(r_i, phi_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sin(theta_M) * z_bar_k / np.sqrt(r_i**2 + z_bar_k**2) - t2 = 1.0 / 2.0 * phi_j * np.cos(phi_bar_M) - t3 = 1.0 / 4.0 * np.sin(phi_bar_M) +def Hr_ri_case212(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sin(theta_M) * z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2) + t2 = 1.0 / 2.0 * phi_j * xp.cos(phi_bar_M) + t3 = 1.0 / 4.0 * xp.sin(phi_bar_M) return t1 * (t2 - t3) -def Hr_phij_case212(r_i, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt(r_i**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hr_phij_case212(xp, r_i, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hphi_ri_case212(r_i, phi_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sin(theta_M) * z_bar_k / np.sqrt(r_i**2 + z_bar_k**2) - t2 = 1.0 / 4.0 * np.cos(phi_bar_M) - t3 = 1.0 / 2.0 * phi_j * np.sin(phi_bar_M) +def Hphi_ri_case212(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sin(theta_M) * z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2) + t2 = 1.0 / 4.0 * xp.cos(phi_bar_M) + t3 = 1.0 / 2.0 * phi_j * xp.sin(phi_bar_M) return t1 * (-t2 + t3) -def Hphi_zk_case212(r_i, theta_M, z_bar_k): - t1 = r_i / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = -np.cos(theta_M) - t2 = np.arctanh(t1) - t2_coef = np.cos(theta_M) +def Hphi_zk_case212(xp, r_i, theta_M, z_bar_k): + t1 = r_i / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = -xp.cos(theta_M) + t2 = xp.atanh(t1) + t2_coef = xp.cos(theta_M) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case212(r_i, phi_bar_M, theta_M, z_bar_k): - t1 = r_i / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_ri_case212(xp, r_i, phi_bar_M, theta_M, z_bar_k): + t1 = r_i / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hz_phij_case212(r_i, phi_bar_M, theta_M, z_bar_k): +def Hz_phij_case212(xp, r_i, phi_bar_M, theta_M, z_bar_k): return ( - np.sin(theta_M) - * np.sin(phi_bar_M) - * np.arctanh(r_i / np.sqrt(r_i**2 + z_bar_k**2)) + xp.sin(theta_M) + * xp.sin(phi_bar_M) + * xp.atanh(r_i / xp.sqrt(r_i**2 + z_bar_k**2)) ) -def Hz_zk_case212(r_i, phi_j, theta_M, z_bar_k): - t1 = phi_j / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = -np.cos(theta_M) * z_bar_k +def Hz_zk_case212(xp, r_i, phi_j, theta_M, z_bar_k): + t1 = phi_j / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = -xp.cos(theta_M) * z_bar_k return t1_coef * t1 # 213 ############## -def Hr_phij_case213(r, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt(r**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hr_phij_case213(xp, r, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt(r**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hphi_zk_case213(r, theta_M, z_bar_k): - t1 = np.sqrt(r**2 + z_bar_k**2) - t1_coef = np.cos(theta_M) / r - t2 = np.arctanh(r / t1) - t2_coef = -np.cos(theta_M) +def Hphi_zk_case213(xp, r, theta_M, z_bar_k): + t1 = xp.sqrt(r**2 + z_bar_k**2) + t1_coef = xp.cos(theta_M) / r + t2 = xp.atanh(r / t1) + t2_coef = -xp.cos(theta_M) return t1_coef * t1 + t2_coef * t2 -def Hz_phij_case213(r, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(r / np.sqrt(r**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_phij_case213(xp, r, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(r / xp.sqrt(r**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hz_zk_case213(phi_bar_j, theta_M, z_bar_k): - t1 = np.sign(z_bar_k) - t1_coef = np.cos(theta_M) * phi_bar_j +def Hz_zk_case213(xp, phi_bar_j, theta_M, z_bar_k): + t1 = xp.sign(z_bar_k) + t1_coef = xp.cos(theta_M) * phi_bar_j return t1_coef * t1 # 214 ############## -def Hr_ri_case214(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): +def Hr_ri_case214(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) E_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k**2 - * np.sign(z_bar_k) + * xp.sign(z_bar_k) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) F_coef = ( - np.sin(theta_M) - * np.cos(phi_bar_M) - * np.sign(z_bar_k) + xp.sin(theta_M) + * xp.cos(phi_bar_M) + * xp.sign(z_bar_k) * (2.0 * r**2 + z_bar_k**2) / (2.0 * r**2) ) return ( - -np.sin(theta_M) - * np.sin(phi_bar_M) - * np.sign(z_bar_k) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) + * xp.sign(z_bar_k) * z_bar_k**2 / (2.0 * r**2) + E_coef * E @@ -485,29 +489,29 @@ def Hr_ri_case214(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): ) -def Hr_phij_case214(phi_bar_M, theta_M, z_bar_k): +def Hr_phij_case214(xp, phi_bar_M, theta_M, z_bar_k): return ( - -np.sin(theta_M) - * np.sin(phi_bar_M) - * np.sign(z_bar_k) - * np.log(np.abs(z_bar_k)) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) + * xp.sign(z_bar_k) + * xp.log(xp.abs(z_bar_k)) ) -def Hr_zk_case214(r, phi_bar_j, theta_M, z_bar_k): +def Hr_zk_case214(xp, r, phi_bar_j, theta_M, z_bar_k): E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) - E_coef = np.cos(theta_M) * np.abs(z_bar_k) / r + E_coef = xp.cos(theta_M) * xp.abs(z_bar_k) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) - F_coef = -np.cos(theta_M) * (2.0 * r**2 + z_bar_k**2) / (r * np.abs(z_bar_k)) - t = np.sqrt(r**2 + z_bar_k**2) + F_coef = -xp.cos(theta_M) * (2.0 * r**2 + z_bar_k**2) / (r * xp.abs(z_bar_k)) + t = xp.sqrt(r**2 + z_bar_k**2) def Pi1(sign): return el3_angle(phi_bar_j / 2, 2 * r / (r + sign * t), -4 * r**2 / z_bar_k**2) def Pi1_coef(sign): return ( - -np.cos(theta_M) - / (r * np.sqrt((r**2 + z_bar_k**2) * z_bar_k**2)) + -xp.cos(theta_M) + / (r * xp.sqrt((r**2 + z_bar_k**2) * z_bar_k**2)) * (t - sign * r) * (r + sign * t) ** 2 ) @@ -522,11 +526,11 @@ def Pi2(sign): def Pi2_coef(sign): return ( sign - * np.cos(theta_M) + * xp.cos(theta_M) * z_bar_k**4 / ( r - * np.sqrt((r**2 + z_bar_k**2) * (4.0 * r**2 + z_bar_k**2)) + * xp.sqrt((r**2 + z_bar_k**2) * (4.0 * r**2 + z_bar_k**2)) * (r + sign * t) ) ) @@ -541,87 +545,87 @@ def Pi2_coef(sign): ) -def Hphi_ri_case214(r, phi_j, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = -np.sin(theta_M) * np.cos(phi_bar_M) * np.sign(z_bar_k) / 2.0 +def Hphi_ri_case214(xp, r, phi_j, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = -xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.sign(z_bar_k) / 2.0 t2 = phi_j - t2_coef = np.sin(theta_M) * np.sin(phi_bar_M) * np.sign(z_bar_k) / 2.0 - t3 = np.sign(z_bar_k) * z_bar_k**2 / (2.0 * r**2) - t3_coef = -np.sin(theta_M) * np.cos(phi_bar_M) - t4 = np.log(np.abs(z_bar_k) / (np.sqrt(2.0) * r)) - t4_coef = -np.sin(theta_M) * np.cos(phi_bar_M) * np.sign(z_bar_k) + t2_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) * xp.sign(z_bar_k) / 2.0 + t3 = xp.sign(z_bar_k) * z_bar_k**2 / (2.0 * r**2) + t3_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) + t4 = xp.log(xp.abs(z_bar_k) / (xp.sqrt(2.0) * r)) + t4_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.sign(z_bar_k) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) E_coef = ( - np.sin(theta_M) - * np.sin(phi_bar_M) + xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k**2 - * np.sign(z_bar_k) + * xp.sign(z_bar_k) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) F_coef = ( - -np.sin(theta_M) - * np.sin(phi_bar_M) - * np.sign(z_bar_k) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) + * xp.sign(z_bar_k) * (4.0 * r**2 + z_bar_k**2) / (2.0 * r**2) ) return t1 + t2_coef * t2 + t3_coef * t3 + t4_coef * t4 + E_coef * E + F_coef * F -def Hphi_zk_case214(r, theta_M, z_bar_k): - t1 = np.abs(z_bar_k) - t1_coef = np.cos(theta_M) / r +def Hphi_zk_case214(xp, r, theta_M, z_bar_k): + t1 = xp.abs(z_bar_k) + t1_coef = xp.cos(theta_M) / r return t1_coef * t1 -def Hz_ri_case214(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): +def Hz_ri_case214(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) - E_coef = np.sin(theta_M) * np.cos(phi_bar_M) * np.abs(z_bar_k) / r + E_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.abs(z_bar_k) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) F_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * (2.0 * r**2 + z_bar_k**2) - / (r * np.abs(z_bar_k)) + / (r * xp.abs(z_bar_k)) ) return ( - np.sin(theta_M) * np.sin(phi_bar_M) * np.abs(z_bar_k) / r + xp.sin(theta_M) * xp.sin(phi_bar_M) * xp.abs(z_bar_k) / r + E_coef * E + F_coef * F ) -def Hz_zk_case214(r, phi_bar_j, theta_M, z_bar_k): - t = np.sqrt(r**2 + z_bar_k**2) +def Hz_zk_case214(xp, r, phi_bar_j, theta_M, z_bar_k): + t = xp.sqrt(r**2 + z_bar_k**2) def Pi(sign): return el3_angle(phi_bar_j / 2, 2 * r / (r + sign * t), -4 * r**2 / z_bar_k**2) - Pi_coef = np.cos(theta_M) * np.sign(z_bar_k) + Pi_coef = xp.cos(theta_M) * xp.sign(z_bar_k) return Pi_coef * Pi(1) + Pi_coef * Pi(-1) # 215 ############## -def Hr_ri_case215(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t2 = np.arctanh(z_bar_k / np.sqrt(r_bar_i**2 + z_bar_k**2)) - t2_coef = np.sin(theta_M) * np.sin(phi_bar_M) / 2.0 * (1.0 - r_i**2 / r**2) +def Hr_ri_case215(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t2 = xp.atanh(z_bar_k / xp.sqrt(r_bar_i**2 + z_bar_k**2)) + t2_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) / 2.0 * (1.0 - r_i**2 / r**2) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) E_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k - * np.sqrt(r_bar_i**2 + z_bar_k**2) + * xp.sqrt(r_bar_i**2 + z_bar_k**2) / (2 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) F_coef = ( - np.sin(theta_M) - * np.cos(phi_bar_M) + xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k * (2.0 * r_i**2 + z_bar_k**2) - / (2 * r**2 * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2 * r**2 * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) Pi = el3_angle( phi_bar_j / 2.0, @@ -629,18 +633,18 @@ def Hr_ri_case215(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2), ) Pi_coef = ( - np.sin(theta_M) - * np.cos(phi_bar_M) + xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k * (r**2 + r_i**2) * (r + r_i) - / (2.0 * r**2 * r_bar_i * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * r_bar_i * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) return ( - -np.sin(theta_M) - * np.sin(phi_bar_M) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k - * np.sqrt(r_bar_i**2 + z_bar_k**2) + * xp.sqrt(r_bar_i**2 + z_bar_k**2) / (2.0 * r**2) + t2_coef * t2 + E_coef * E @@ -649,22 +653,22 @@ def Hr_ri_case215(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): ) -def Hr_phij_case215(r_bar_i, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt(r_bar_i**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hr_phij_case215(xp, r_bar_i, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt(r_bar_i**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hr_zk_case215(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): +def Hr_zk_case215(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) - E_coef = np.cos(theta_M) * np.sqrt(r_bar_i**2 + z_bar_k**2) / r + E_coef = xp.cos(theta_M) * xp.sqrt(r_bar_i**2 + z_bar_k**2) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) F_coef = ( - -np.cos(theta_M) + -xp.cos(theta_M) * (r**2 + r_i**2 + z_bar_k**2) - / (r * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (r * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) - t = np.sqrt(r**2 + z_bar_k**2) + t = xp.sqrt(r**2 + z_bar_k**2) def Pi1(sign): return el3_angle( @@ -675,8 +679,8 @@ def Pi1(sign): def Pi1_coef(sign): return ( - -np.cos(theta_M) - / (r * np.sqrt((r**2 + z_bar_k**2) * (r_bar_i**2 + z_bar_k**2))) + -xp.cos(theta_M) + / (r * xp.sqrt((r**2 + z_bar_k**2) * (r_bar_i**2 + z_bar_k**2))) * (t - sign * r) * (r_i + sign * t) ** 2 ) @@ -694,12 +698,12 @@ def Pi2(sign): def Pi2_coef(sign): return ( sign - * np.cos(theta_M) + * xp.cos(theta_M) * z_bar_k**2 * (r_bar_i**2 + z_bar_k**2) / ( r - * np.sqrt((r**2 + z_bar_k**2) * ((r + r_i) ** 2 + z_bar_k**2)) + * xp.sqrt((r**2 + z_bar_k**2) * ((r + r_i) ** 2 + z_bar_k**2)) * (r + sign * t) ) ) @@ -714,26 +718,26 @@ def Pi2_coef(sign): ) -def Hphi_ri_case215(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt(r_bar_i**2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) - t1_coef = -np.sin(theta_M) * np.cos(phi_bar_M) - t2 = np.arctanh(z_bar_k / np.sqrt(r_bar_i**2 + z_bar_k**2)) - t2_coef = -np.sin(theta_M) * np.cos(phi_bar_M) * (r**2 + r_i**2) / (2.0 * r**2) +def Hphi_ri_case215(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt(r_bar_i**2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) + t1_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) + t2 = xp.atanh(z_bar_k / xp.sqrt(r_bar_i**2 + z_bar_k**2)) + t2_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) * (r**2 + r_i**2) / (2.0 * r**2) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) E_coef = ( - np.sin(theta_M) - * np.sin(phi_bar_M) + xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k - * np.sqrt(r_bar_i**2 + z_bar_k**2) + * xp.sqrt(r_bar_i**2 + z_bar_k**2) / (2 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) F_coef = ( - -np.sin(theta_M) - * np.sin(phi_bar_M) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k * (2.0 * r**2 + 2.0 * r_i**2 + z_bar_k**2) - / (2.0 * r**2 * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) Pi = el3_angle( phi_bar_j / 2.0, @@ -741,47 +745,47 @@ def Hphi_ri_case215(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2), ) Pi_coef = ( - np.sin(theta_M) - * np.sin(phi_bar_M) + xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k * (r + r_i) ** 2 - / (2.0 * r**2 * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) return t1_coef * t1 + t2_coef * t2 + E_coef * E + F_coef * F + Pi_coef * Pi -def Hphi_zk_case215(r, r_bar_i, theta_M, z_bar_k): - t1 = np.sqrt(r_bar_i**2 + z_bar_k**2) - t1_coef = np.cos(theta_M) / r - t2 = np.arctanh(r_bar_i / t1) - t2_coef = -np.cos(theta_M) +def Hphi_zk_case215(xp, r, r_bar_i, theta_M, z_bar_k): + t1 = xp.sqrt(r_bar_i**2 + z_bar_k**2) + t1_coef = xp.cos(theta_M) / r + t2 = xp.atanh(r_bar_i / t1) + t2_coef = -xp.cos(theta_M) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case215(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): +def Hz_ri_case215(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): t = r_bar_i**2 + z_bar_k**2 - t1 = np.sqrt(r_bar_i**2 + z_bar_k**2) / r - t1_coef = np.sin(theta_M) * np.sin(phi_bar_M) + t1 = xp.sqrt(r_bar_i**2 + z_bar_k**2) / r + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / t) - E_coef = np.sin(theta_M) * np.cos(phi_bar_M) * np.sqrt(t) / r + E_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.sqrt(t) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / t) F_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * (r**2 + r_i**2 + z_bar_k**2) - / (r * np.sqrt(t)) + / (r * xp.sqrt(t)) ) return t1_coef * t1 + E_coef * E + F_coef * F -def Hz_phij_case215(r_bar_i, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(r_bar_i / np.sqrt(r_bar_i**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_phij_case215(xp, r_bar_i, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(r_bar_i / xp.sqrt(r_bar_i**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hz_zk_case215(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): - t = np.sqrt(r**2 + z_bar_k**2) +def Hz_zk_case215(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): + t = xp.sqrt(r**2 + z_bar_k**2) def Pi(sign): return el3_angle( @@ -792,10 +796,10 @@ def Pi(sign): def Pi_coef(sign): return ( - np.cos(theta_M) + xp.cos(theta_M) * z_bar_k * (r_i + sign * t) - / (np.sqrt(r_bar_i**2 + z_bar_k**2) * (r + sign * t)) + / (xp.sqrt(r_bar_i**2 + z_bar_k**2) * (r + sign * t)) ) return Pi_coef(1) * Pi(1) + Pi_coef(-1) * Pi(-1) @@ -804,134 +808,134 @@ def Pi_coef(sign): # 221 ############## -def Hr_phij_case221(phi_bar_M, theta_M, z_bar_k): +def Hr_phij_case221(xp, phi_bar_M, theta_M, z_bar_k): return ( - -np.sin(theta_M) - * np.sin(phi_bar_M) - * np.sign(z_bar_k) - * np.log(np.abs(z_bar_k)) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) + * xp.sign(z_bar_k) + * xp.log(xp.abs(z_bar_k)) ) -def Hz_zk_case221(phi_j, theta_M, z_bar_k): - return -np.cos(theta_M) * np.sign(z_bar_k) * phi_j +def Hz_zk_case221(xp, phi_j, theta_M, z_bar_k): + return -xp.cos(theta_M) * xp.sign(z_bar_k) * phi_j # 222 ############## -def Hr_ri_case222(r_i, phi_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sin(theta_M) * z_bar_k / np.sqrt(r_i**2 + z_bar_k**2) - t2 = 1.0 / 2.0 * phi_j * np.cos(phi_bar_M) - t3 = 1.0 / 4.0 * np.sin(phi_bar_M) +def Hr_ri_case222(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sin(theta_M) * z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2) + t2 = 1.0 / 2.0 * phi_j * xp.cos(phi_bar_M) + t3 = 1.0 / 4.0 * xp.sin(phi_bar_M) return t1 * (t2 - t3) -def Hr_phij_case222(r_i, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt(r_i**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hr_phij_case222(xp, r_i, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hphi_ri_case222(r_i, phi_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sin(theta_M) * z_bar_k / np.sqrt(r_i**2 + z_bar_k**2) - t2 = 1.0 / 4.0 * np.cos(phi_bar_M) - t3 = 1.0 / 2.0 * phi_j * np.sin(phi_bar_M) +def Hphi_ri_case222(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sin(theta_M) * z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2) + t2 = 1.0 / 4.0 * xp.cos(phi_bar_M) + t3 = 1.0 / 2.0 * phi_j * xp.sin(phi_bar_M) return t1 * (-t2 + t3) -def Hphi_zk_case222(r_i, theta_M, z_bar_k): - t1 = r_i / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = np.cos(theta_M) - t2 = np.arctanh(t1) - t2_coef = -np.cos(theta_M) +def Hphi_zk_case222(xp, r_i, theta_M, z_bar_k): + t1 = r_i / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = xp.cos(theta_M) + t2 = xp.atanh(t1) + t2_coef = -xp.cos(theta_M) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case222(r_i, phi_bar_M, theta_M, z_bar_k): - t1 = r_i / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_ri_case222(xp, r_i, phi_bar_M, theta_M, z_bar_k): + t1 = r_i / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hz_phij_case222(r_i, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(r_i / np.sqrt(r_i**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_phij_case222(xp, r_i, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(r_i / xp.sqrt(r_i**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hz_zk_case222(r_i, phi_j, theta_M, z_bar_k): - t1 = z_bar_k / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = -np.cos(theta_M) * phi_j +def Hz_zk_case222(xp, r_i, phi_j, theta_M, z_bar_k): + t1 = z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = -xp.cos(theta_M) * phi_j return t1_coef * t1 # 223 ############## -def Hr_phij_case223(r, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt(r**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hr_phij_case223(xp, r, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt(r**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hphi_zk_case223(r, theta_M, z_bar_k): - t1 = np.sqrt(r**2 + z_bar_k**2) - t1_coef = np.cos(theta_M) / r - t2 = np.arctanh(r / t1) - t2_coef = -np.cos(theta_M) +def Hphi_zk_case223(xp, r, theta_M, z_bar_k): + t1 = xp.sqrt(r**2 + z_bar_k**2) + t1_coef = xp.cos(theta_M) / r + t2 = xp.atanh(r / t1) + t2_coef = -xp.cos(theta_M) return t1_coef * t1 + t2_coef * t2 -def Hz_phij_case223(r, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(r / np.sqrt(r**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_phij_case223(xp, r, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(r / xp.sqrt(r**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hz_zk_case223(r, phi_bar_j, theta_M, z_bar_k): - t1 = arctan_k_tan_2(np.sqrt(r**2 + z_bar_k**2) / np.abs(z_bar_k), 2.0 * phi_bar_j) - t1_coef = np.cos(theta_M) * np.sign(z_bar_k) +def Hz_zk_case223(xp, r, phi_bar_j, theta_M, z_bar_k): + t1 = arctan_k_tan_2(xp.sqrt(r**2 + z_bar_k**2) / xp.abs(z_bar_k), 2.0 * phi_bar_j) + t1_coef = xp.cos(theta_M) * xp.sign(z_bar_k) return t1_coef * t1 # 224 ############## -def Hr_ri_case224(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt(4.0 * r**2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hr_ri_case224(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt(4.0 * r**2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) E_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k**2 - * np.sign(z_bar_k) + * xp.sign(z_bar_k) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) F_coef = ( - np.sin(theta_M) - * np.cos(phi_bar_M) - * np.sign(z_bar_k) + xp.sin(theta_M) + * xp.cos(phi_bar_M) + * xp.sign(z_bar_k) * (2.0 * r**2 + z_bar_k**2) / (2.0 * r**2) ) return t1_coef * t1 + E_coef * E + F_coef * F -def Hr_phij_case224(r, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt(4.0 * r**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hr_phij_case224(xp, r, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt(4.0 * r**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hr_zk_case224(r, phi_bar_j, theta_M, z_bar_k): +def Hr_zk_case224(xp, r, phi_bar_j, theta_M, z_bar_k): E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) - E_coef = np.cos(theta_M) * np.abs(z_bar_k) / r + E_coef = xp.cos(theta_M) * xp.abs(z_bar_k) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) - F_coef = -np.cos(theta_M) * (2.0 * r**2 + z_bar_k**2) / (r * np.abs(z_bar_k)) - t = np.sqrt(r**2 + z_bar_k**2) + F_coef = -xp.cos(theta_M) * (2.0 * r**2 + z_bar_k**2) / (r * xp.abs(z_bar_k)) + t = xp.sqrt(r**2 + z_bar_k**2) def Pi1(sign): return el3_angle( @@ -940,8 +944,8 @@ def Pi1(sign): def Pi1_coef(sign): return ( - -np.cos(theta_M) - / (r * np.sqrt((r**2 + z_bar_k**2) * z_bar_k**2)) + -xp.cos(theta_M) + / (r * xp.sqrt((r**2 + z_bar_k**2) * z_bar_k**2)) * (t - sign * r) * (r + sign * t) ** 2 ) @@ -956,11 +960,11 @@ def Pi2(sign): def Pi2_coef(sign): return ( sign - * np.cos(theta_M) + * xp.cos(theta_M) * z_bar_k**4 / ( r - * np.sqrt((r**2 + z_bar_k**2) * (4.0 * r**2 + z_bar_k**2)) + * xp.sqrt((r**2 + z_bar_k**2) * (4.0 * r**2 + z_bar_k**2)) * (r + sign * t) ) ) @@ -975,94 +979,94 @@ def Pi2_coef(sign): ) -def Hphi_ri_case224(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt(4.0 * r**2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) - t1_coef = -np.sin(theta_M) * np.cos(phi_bar_M) - t2 = np.arctanh(z_bar_k / np.sqrt(4.0 * r**2 + z_bar_k**2)) - t2_coef = -np.sin(theta_M) * np.cos(phi_bar_M) +def Hphi_ri_case224(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt(4.0 * r**2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) + t1_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) + t2 = xp.atanh(z_bar_k / xp.sqrt(4.0 * r**2 + z_bar_k**2)) + t2_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) E_coef = ( - np.sin(theta_M) - * np.sin(phi_bar_M) + xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k**2 - * np.sign(z_bar_k) + * xp.sign(z_bar_k) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) F_coef = ( - -np.sin(theta_M) - * np.sin(phi_bar_M) - * np.sign(z_bar_k) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) + * xp.sign(z_bar_k) * (4.0 * r**2 + z_bar_k**2) / (2.0 * r**2) ) return t1_coef * t1 + t2_coef * t2 + E_coef * E + F_coef * F -def Hphi_zk_case224(r, theta_M, z_bar_k): - t1 = np.sqrt(4.0 * r**2 + z_bar_k**2) - t1_coef = np.cos(theta_M) / r - t2 = np.arctanh(2.0 * r / t1) - t2_coef = -np.cos(theta_M) +def Hphi_zk_case224(xp, r, theta_M, z_bar_k): + t1 = xp.sqrt(4.0 * r**2 + z_bar_k**2) + t1_coef = xp.cos(theta_M) / r + t2 = xp.atanh(2.0 * r / t1) + t2_coef = -xp.cos(theta_M) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case224(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt(4.0 * r**2 + z_bar_k**2) / r - t1_coef = np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_ri_case224(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt(4.0 * r**2 + z_bar_k**2) / r + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) - E_coef = np.sin(theta_M) * np.cos(phi_bar_M) * np.abs(z_bar_k) / r + E_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.abs(z_bar_k) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) F_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * (2.0 * r**2 + z_bar_k**2) - / (r * np.abs(z_bar_k)) + / (r * xp.abs(z_bar_k)) ) return t1_coef * t1 + E_coef * E + F_coef * F -def Hz_phij_case224(r, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(2.0 * r / np.sqrt(4.0 * r**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_phij_case224(xp, r, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(2.0 * r / xp.sqrt(4.0 * r**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hz_zk_case224(r, phi_bar_j, theta_M, z_bar_k): - t = np.sqrt(r**2 + z_bar_k**2) +def Hz_zk_case224(xp, r, phi_bar_j, theta_M, z_bar_k): + t = xp.sqrt(r**2 + z_bar_k**2) def Pi(sign): return el3_angle( phi_bar_j / 2.0, 2.0 * r / (r + sign * t), -4.0 * r**2 / z_bar_k**2 ) - Pi_coef = np.cos(theta_M) * np.sign(z_bar_k) + Pi_coef = xp.cos(theta_M) * xp.sign(z_bar_k) return Pi_coef * Pi(1) + Pi_coef * Pi(-1) # 225 ############## -def Hr_ri_case225(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt((r + r_i) ** 2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) - t2 = np.arctanh(z_bar_k / np.sqrt((r + r_i) ** 2 + z_bar_k**2)) - t2_coef = np.sin(theta_M) * np.sin(phi_bar_M) / 2.0 * (1.0 - r_i**2 / r**2) +def Hr_ri_case225(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt((r + r_i) ** 2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) + t2 = xp.atanh(z_bar_k / xp.sqrt((r + r_i) ** 2 + z_bar_k**2)) + t2_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) / 2.0 * (1.0 - r_i**2 / r**2) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) E_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k - * np.sqrt(r_bar_i**2 + z_bar_k**2) + * xp.sqrt(r_bar_i**2 + z_bar_k**2) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) F_coef = ( - np.sin(theta_M) - * np.cos(phi_bar_M) + xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k * (2.0 * r_i**2 + z_bar_k**2) - / (2.0 * r**2 * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) Pi = el3_angle( phi_bar_j / 2.0, @@ -1070,32 +1074,32 @@ def Hr_ri_case225(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2), ) Pi_coef = ( - np.sin(theta_M) - * np.cos(phi_bar_M) + xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k * (r**2 + r_i**2) * (r + r_i) - / (2.0 * r**2 * r_bar_i * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * r_bar_i * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) return t1_coef * t1 + t2_coef * t2 + E_coef * E + F_coef * F + Pi_coef * Pi -def Hr_phij_case225(r, r_i, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt((r + r_i) ** 2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hr_phij_case225(xp, r, r_i, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt((r + r_i) ** 2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hr_zk_case225(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): +def Hr_zk_case225(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) - E_coef = np.cos(theta_M) * np.sqrt(r_bar_i**2 + z_bar_k**2) / r + E_coef = xp.cos(theta_M) * xp.sqrt(r_bar_i**2 + z_bar_k**2) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) F_coef = ( - -np.cos(theta_M) + -xp.cos(theta_M) * (r**2 + r_i**2 + z_bar_k**2) - / (r * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (r * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) - t = np.sqrt(r**2 + z_bar_k**2) + t = xp.sqrt(r**2 + z_bar_k**2) def Pi1(sign): return el3_angle( @@ -1106,8 +1110,8 @@ def Pi1(sign): def Pi1_coef(sign): return ( - -np.cos(theta_M) - / (r * np.sqrt((r**2 + z_bar_k**2) * (r_bar_i**2 + z_bar_k**2))) + -xp.cos(theta_M) + / (r * xp.sqrt((r**2 + z_bar_k**2) * (r_bar_i**2 + z_bar_k**2))) * (t - sign * r) * (r_i + sign * t) ** 2 ) @@ -1125,12 +1129,12 @@ def Pi2(sign): def Pi2_coef(sign): return ( sign - * np.cos(theta_M) + * xp.cos(theta_M) * z_bar_k**2 * (r_bar_i**2 + z_bar_k**2) / ( r - * np.sqrt((r**2 + z_bar_k**2) * ((r + r_i) ** 2 + z_bar_k**2)) + * xp.sqrt((r**2 + z_bar_k**2) * ((r + r_i) ** 2 + z_bar_k**2)) * (r + sign * t) ) ) @@ -1145,26 +1149,26 @@ def Pi2_coef(sign): ) -def Hphi_ri_case225(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt((r + r_i) ** 2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) - t1_coef = -np.sin(theta_M) * np.cos(phi_bar_M) - t2 = np.arctanh(z_bar_k / np.sqrt((r + r_i) ** 2 + z_bar_k**2)) - t2_coef = -np.sin(theta_M) * np.cos(phi_bar_M) * (r**2 + r_i**2) / (2.0 * r**2) +def Hphi_ri_case225(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt((r + r_i) ** 2 + z_bar_k**2) * z_bar_k / (2.0 * r**2) + t1_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) + t2 = xp.atanh(z_bar_k / xp.sqrt((r + r_i) ** 2 + z_bar_k**2)) + t2_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) * (r**2 + r_i**2) / (2.0 * r**2) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) E_coef = ( - np.sin(theta_M) - * np.sin(phi_bar_M) + xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k - * np.sqrt(r_bar_i**2 + z_bar_k**2) + * xp.sqrt(r_bar_i**2 + z_bar_k**2) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) F_coef = ( - -np.sin(theta_M) - * np.sin(phi_bar_M) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k * (2.0 * r**2 + 2.0 * r_i**2 + z_bar_k**2) - / (2.0 * r**2 * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) Pi = el3_angle( phi_bar_j / 2.0, @@ -1172,47 +1176,47 @@ def Hphi_ri_case225(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2), ) Pi_coef = ( - np.sin(theta_M) - * np.sin(phi_bar_M) + xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k * (r + r_i) ** 2 - / (2.0 * r**2 * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) return t1_coef * t1 + t2_coef * t2 + E_coef * E + F_coef * F + Pi_coef * Pi -def Hphi_zk_case225(r, r_i, theta_M, z_bar_k): - t1 = np.sqrt((r + r_i) ** 2 + z_bar_k**2) - t1_coef = np.cos(theta_M) / r - t2 = np.arctanh((r + r_i) / t1) - t2_coef = -np.cos(theta_M) +def Hphi_zk_case225(xp, r, r_i, theta_M, z_bar_k): + t1 = xp.sqrt((r + r_i) ** 2 + z_bar_k**2) + t1_coef = xp.cos(theta_M) / r + t2 = xp.atanh((r + r_i) / t1) + t2_coef = -xp.cos(theta_M) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case225(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): +def Hz_ri_case225(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): t = r_bar_i**2 + z_bar_k**2 - t1 = np.sqrt((r + r_i) ** 2 + z_bar_k**2) / r - t1_coef = np.sin(theta_M) * np.sin(phi_bar_M) + t1 = xp.sqrt((r + r_i) ** 2 + z_bar_k**2) / r + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / t) - E_coef = np.sin(theta_M) * np.cos(phi_bar_M) * np.sqrt(t) / r + E_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.sqrt(t) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / t) F_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * (r**2 + r_i**2 + z_bar_k**2) - / (r * np.sqrt(t)) + / (r * xp.sqrt(t)) ) return t1_coef * t1 + E_coef * E + F_coef * F -def Hz_phij_case225(r, r_i, phi_bar_M, theta_M, z_bar_k): - t1 = np.arctanh((r + r_i) / np.sqrt((r + r_i) ** 2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) +def Hz_phij_case225(xp, r, r_i, phi_bar_M, theta_M, z_bar_k): + t1 = xp.atanh((r + r_i) / xp.sqrt((r + r_i) ** 2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) return t1_coef * t1 -def Hz_zk_case225(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): - t = np.sqrt(r**2 + z_bar_k**2) +def Hz_zk_case225(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): + t = xp.sqrt(r**2 + z_bar_k**2) def Pi(sign): return el3_angle( @@ -1223,10 +1227,10 @@ def Pi(sign): def Pi_coef(sign): return ( - np.cos(theta_M) + xp.cos(theta_M) * z_bar_k * (r_i + sign * t) - / (np.sqrt(r_bar_i**2 + z_bar_k**2) * (r + sign * t)) + / (xp.sqrt(r_bar_i**2 + z_bar_k**2) * (r + sign * t)) ) return Pi_coef(1) * Pi(1) + Pi_coef(-1) * Pi(-1) @@ -1235,119 +1239,119 @@ def Pi_coef(sign): # 231 ############## -def Hr_phij_case231(phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): +def Hr_phij_case231(xp, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): return ( - -np.sin(theta_M) - * np.sin(phi_bar_Mj) - * np.cos(phi_bar_j) - * np.sign(z_bar_k) - * np.log(np.abs(z_bar_k)) + -xp.sin(theta_M) + * xp.sin(phi_bar_Mj) + * xp.cos(phi_bar_j) + * xp.sign(z_bar_k) + * xp.log(xp.abs(z_bar_k)) ) -def Hphi_phij_case231(phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.log(np.abs(z_bar_k)) +def Hphi_phij_case231(xp, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.log(xp.abs(z_bar_k)) t1_coef = ( - np.sin(theta_M) * np.sin(phi_bar_Mj) * np.sin(phi_bar_j) * np.sign(z_bar_k) + xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.sin(phi_bar_j) * xp.sign(z_bar_k) ) return t1_coef * t1 -def Hz_zk_case231(phi_j, theta_M, z_bar_k): - t1 = phi_j * np.sign(z_bar_k) - t1_coef = -np.cos(theta_M) +def Hz_zk_case231(xp, phi_j, theta_M, z_bar_k): + t1 = phi_j * xp.sign(z_bar_k) + t1_coef = -xp.cos(theta_M) return t1_coef * t1 # 232 ############## -def Hr_ri_case232(r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.sin(theta_M) * z_bar_k / np.sqrt(r_i**2 + z_bar_k**2) - t2 = 1.0 / 2.0 * phi_j * np.cos(phi_bar_M) - t3 = 1.0 / 4.0 * np.sin(phi_bar_Mj + phi_bar_j) +def Hr_ri_case232(xp, r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.sin(theta_M) * z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2) + t2 = 1.0 / 2.0 * phi_j * xp.cos(phi_bar_M) + t3 = 1.0 / 4.0 * xp.sin(phi_bar_Mj + phi_bar_j) return t1 * (t2 - t3) -def Hr_phij_case232(r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt(r_i**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) * np.cos(phi_bar_j) +def Hr_phij_case232(xp, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.cos(phi_bar_j) return t1_coef * t1 -def Hr_zk_case232(r_i, phi_bar_j, theta_M, z_bar_k): - t1 = r_i / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = -np.cos(theta_M) * np.sin(phi_bar_j) - t2 = np.arctanh(t1) - t2_coef = np.cos(theta_M) * np.sin(phi_bar_j) +def Hr_zk_case232(xp, r_i, phi_bar_j, theta_M, z_bar_k): + t1 = r_i / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = -xp.cos(theta_M) * xp.sin(phi_bar_j) + t2 = xp.atanh(t1) + t2_coef = xp.cos(theta_M) * xp.sin(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hphi_ri_case232(r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.sin(theta_M) * z_bar_k / np.sqrt(r_i**2 + z_bar_k**2) - t2 = 1.0 / 4.0 * np.cos(phi_bar_Mj + phi_bar_j) - t3 = 1.0 / 2.0 * phi_j * np.sin(phi_bar_M) +def Hphi_ri_case232(xp, r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.sin(theta_M) * z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2) + t2 = 1.0 / 4.0 * xp.cos(phi_bar_Mj + phi_bar_j) + t3 = 1.0 / 2.0 * phi_j * xp.sin(phi_bar_M) return t1 * (-t2 + t3) -def Hphi_phij_case232(r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt(r_i**2 + z_bar_k**2)) - t1_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) * np.sin(phi_bar_j) +def Hphi_phij_case232(xp, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2)) + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.sin(phi_bar_j) return t1_coef * t1 -def Hphi_zk_case232(r_i, phi_bar_j, theta_M, z_bar_k): - t1 = r_i / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = -np.cos(theta_M) * np.cos(phi_bar_j) - t2 = np.arctanh(t1) - t2_coef = np.cos(theta_M) * np.cos(phi_bar_j) +def Hphi_zk_case232(xp, r_i, phi_bar_j, theta_M, z_bar_k): + t1 = r_i / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = -xp.cos(theta_M) * xp.cos(phi_bar_j) + t2 = xp.atanh(t1) + t2_coef = xp.cos(theta_M) * xp.cos(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case232(r_i, phi_bar_Mj, theta_M, z_bar_k): - t1 = r_i / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) +def Hz_ri_case232(xp, r_i, phi_bar_Mj, theta_M, z_bar_k): + t1 = r_i / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) return t1_coef * t1 -def Hz_phij_case232(r_i, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.arctanh(r_i / np.sqrt(r_i**2 + z_bar_k**2)) - t1_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) +def Hz_phij_case232(xp, r_i, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.atanh(r_i / xp.sqrt(r_i**2 + z_bar_k**2)) + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) return t1_coef * t1 -def Hz_zk_case232(r_i, phi_j, theta_M, z_bar_k): - t1 = z_bar_k / np.sqrt(r_i**2 + z_bar_k**2) - t1_coef = -np.cos(theta_M) * phi_j +def Hz_zk_case232(xp, r_i, phi_j, theta_M, z_bar_k): + t1 = z_bar_k / xp.sqrt(r_i**2 + z_bar_k**2) + t1_coef = -xp.cos(theta_M) * phi_j return t1_coef * t1 # 233 ############## -def Hr_phij_case233(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.arctanh(z_bar_k / np.sqrt(r**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) * np.cos(phi_bar_j) - t2 = np.arctan( - z_bar_k * np.cos(phi_bar_j) / np.sin(phi_bar_j) / np.sqrt(r**2 + z_bar_k**2) +def Hr_phij_case233(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.atanh(z_bar_k / xp.sqrt(r**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.cos(phi_bar_j) + t2 = xp.atan( + z_bar_k * xp.cos(phi_bar_j) / xp.sin(phi_bar_j) / xp.sqrt(r**2 + z_bar_k**2) ) - t2_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) * np.sin(phi_bar_j) + t2_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.sin(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hr_zk_case233(r, phi_bar_j, theta_M, z_bar_k): - t = np.sqrt(r**2 + z_bar_k**2) - t1 = np.sin(phi_bar_j) - t1_coef = -np.cos(theta_M) - t2 = np.log(-r * np.cos(phi_bar_j) + t) - t2_coef = np.cos(theta_M) * np.sin(phi_bar_j) - t3 = np.arctan(r * np.sin(phi_bar_j) / z_bar_k) - t3_coef = np.cos(theta_M) * z_bar_k / r - t4 = arctan_k_tan_2(t / np.abs(z_bar_k), 2.0 * phi_bar_j) +def Hr_zk_case233(xp, r, phi_bar_j, theta_M, z_bar_k): + t = xp.sqrt(r**2 + z_bar_k**2) + t1 = xp.sin(phi_bar_j) + t1_coef = -xp.cos(theta_M) + t2 = xp.log(-r * xp.cos(phi_bar_j) + t) + t2_coef = xp.cos(theta_M) * xp.sin(phi_bar_j) + t3 = xp.atan(r * xp.sin(phi_bar_j) / z_bar_k) + t3_coef = xp.cos(theta_M) * z_bar_k / r + t4 = arctan_k_tan_2(t / xp.abs(z_bar_k), 2.0 * phi_bar_j) t4_coef = -t3_coef def t5(sign): - return arctan_k_tan_2(np.abs(z_bar_k) / np.abs(r + sign * t), phi_bar_j) + return arctan_k_tan_2(xp.abs(z_bar_k) / xp.abs(r + sign * t), phi_bar_j) t5_coef = t3_coef return ( @@ -1360,92 +1364,92 @@ def t5(sign): ) -def Hphi_phij_case233(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t = np.sqrt(r**2 + z_bar_k**2) - t1 = np.arctan(z_bar_k * np.cos(phi_bar_j) / (np.sin(phi_bar_j) * t)) - t1_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) * np.cos(phi_bar_j) - t2 = np.arctanh(z_bar_k / t) - t2_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) * np.sin(phi_bar_j) +def Hphi_phij_case233(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t = xp.sqrt(r**2 + z_bar_k**2) + t1 = xp.atan(z_bar_k * xp.cos(phi_bar_j) / (xp.sin(phi_bar_j) * t)) + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.cos(phi_bar_j) + t2 = xp.atanh(z_bar_k / t) + t2_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.sin(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hphi_zk_case233(r, phi_bar_j, theta_M, z_bar_k): - t1 = np.sqrt(r**2 + z_bar_k**2) - t1_coef = np.cos(theta_M) / r - t2 = np.arctanh(r * np.cos(phi_bar_j) / t1) - t2_coef = -np.cos(theta_M) * np.cos(phi_bar_j) +def Hphi_zk_case233(xp, r, phi_bar_j, theta_M, z_bar_k): + t1 = xp.sqrt(r**2 + z_bar_k**2) + t1_coef = xp.cos(theta_M) / r + t2 = xp.atanh(r * xp.cos(phi_bar_j) / t1) + t2_coef = -xp.cos(theta_M) * xp.cos(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hz_phij_case233(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.arctanh(r * np.cos(phi_bar_j) / np.sqrt(r**2 + z_bar_k**2)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) +def Hz_phij_case233(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.atanh(r * xp.cos(phi_bar_j) / xp.sqrt(r**2 + z_bar_k**2)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) return t1_coef * t1 -def Hz_zk_case233(r, phi_bar_j, theta_M, z_bar_k): - t1 = arctan_k_tan_2(np.sqrt(r**2 + z_bar_k**2) / np.abs(z_bar_k), 2.0 * phi_bar_j) - t1_coef = np.cos(theta_M) * np.sign(z_bar_k) +def Hz_zk_case233(xp, r, phi_bar_j, theta_M, z_bar_k): + t1 = arctan_k_tan_2(xp.sqrt(r**2 + z_bar_k**2) / xp.abs(z_bar_k), 2.0 * phi_bar_j) + t1_coef = xp.cos(theta_M) * xp.sign(z_bar_k) return t1_coef * t1 # 234 ############## -def Hr_ri_case234(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) * z_bar_k / (2.0 * r**2) +def Hr_ri_case234(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) * z_bar_k / (2.0 * r**2) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) E_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k**2 - * np.sign(z_bar_k) + * xp.sign(z_bar_k) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) F_coef = ( - np.sin(theta_M) - * np.cos(phi_bar_M) - * np.sign(z_bar_k) + xp.sin(theta_M) + * xp.cos(phi_bar_M) + * xp.sign(z_bar_k) * (2.0 * r**2 + z_bar_k**2) / (2.0 * r**2) ) return t1_coef * t1 + E_coef * E + F_coef * F -def Hr_phij_case234(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.arctanh( - z_bar_k / np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) +def Hr_phij_case234(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.atanh( + z_bar_k / xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) ) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) * np.cos(phi_bar_j) - t2 = np.arctan( + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.cos(phi_bar_j) + t2 = xp.atan( z_bar_k - * (1.0 - np.cos(phi_bar_j)) + * (1.0 - xp.cos(phi_bar_j)) / ( - np.sin(phi_bar_j) - * np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) + xp.sin(phi_bar_j) + * xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) ) ) - t2_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) * np.sin(phi_bar_j) + t2_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.sin(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hr_zk_case234(r, phi_bar_j, theta_M, z_bar_k): - t1 = np.sin(phi_bar_j) - t1_coef = -np.cos(theta_M) - t2 = np.log( - r * (1.0 - np.cos(phi_bar_j)) - + np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) +def Hr_zk_case234(xp, r, phi_bar_j, theta_M, z_bar_k): + t1 = xp.sin(phi_bar_j) + t1_coef = -xp.cos(theta_M) + t2 = xp.log( + r * (1.0 - xp.cos(phi_bar_j)) + + xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) ) - t2_coef = np.cos(theta_M) * np.sin(phi_bar_j) - t3 = np.arctan(r * np.sin(phi_bar_j) / z_bar_k) - t3_coef = np.cos(theta_M) * z_bar_k / r + t2_coef = xp.cos(theta_M) * xp.sin(phi_bar_j) + t3 = xp.atan(r * xp.sin(phi_bar_j) / z_bar_k) + t3_coef = xp.cos(theta_M) * z_bar_k / r E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) - E_coef = np.cos(theta_M) * np.abs(z_bar_k) / r + E_coef = xp.cos(theta_M) * xp.abs(z_bar_k) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) - F_coef = -np.cos(theta_M) * (2.0 * r**2 + z_bar_k**2) / (r * np.abs(z_bar_k)) - t = np.sqrt(r**2 + z_bar_k**2) + F_coef = -xp.cos(theta_M) * (2.0 * r**2 + z_bar_k**2) / (r * xp.abs(z_bar_k)) + t = xp.sqrt(r**2 + z_bar_k**2) def Pi1(sign): return el3_angle( @@ -1454,15 +1458,15 @@ def Pi1(sign): def Pi1_coef(sign): return ( - -np.cos(theta_M) - / (r * np.sqrt((r**2 + z_bar_k**2) * z_bar_k**2)) + -xp.cos(theta_M) + / (r * xp.sqrt((r**2 + z_bar_k**2) * z_bar_k**2)) * (t - sign * r) * (r + sign * t) ** 2 ) def Pi2(sign): return el3_angle( - arctan_k_tan_2(np.sqrt((4.0 * r**2 + z_bar_k**2) / z_bar_k**2), phi_bar_j), + arctan_k_tan_2(xp.sqrt((4.0 * r**2 + z_bar_k**2) / z_bar_k**2), phi_bar_j), 1.0 - z_bar_k**4 / ((4.0 * r**2 + z_bar_k**2) * (r + sign * t) ** 2), 4.0 * r**2 / (4.0 * r**2 + z_bar_k**2), ) @@ -1470,11 +1474,11 @@ def Pi2(sign): def Pi2_coef(sign): return ( sign - * np.cos(theta_M) + * xp.cos(theta_M) * z_bar_k**4 / ( r - * np.sqrt((r**2 + z_bar_k**2) * (4.0 * r**2 + z_bar_k**2)) + * xp.sqrt((r**2 + z_bar_k**2) * (4.0 * r**2 + z_bar_k**2)) * (r + sign * t) ) ) @@ -1492,112 +1496,112 @@ def Pi2_coef(sign): ) -def Hphi_ri_case234(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) - t1_coef = -np.sin(theta_M) * np.cos(phi_bar_M) * z_bar_k / (2.0 * r**2) - t2 = np.arctanh( - z_bar_k / np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) +def Hphi_ri_case234(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) + t1_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) * z_bar_k / (2.0 * r**2) + t2 = xp.atanh( + z_bar_k / xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) ) - t2_coef = -np.sin(theta_M) * np.cos(phi_bar_M) + t2_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) E_coef = ( - np.sin(theta_M) - * np.sin(phi_bar_M) + xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k**2 - * np.sign(z_bar_k) + * xp.sign(z_bar_k) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) F_coef = ( - -np.sin(theta_M) - * np.sin(phi_bar_M) - * np.sign(z_bar_k) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) + * xp.sign(z_bar_k) * (4.0 * r**2 + z_bar_k**2) / (2.0 * r**2) ) return t1_coef * t1 + t2_coef * t2 + E_coef * E + F_coef * F -def Hphi_phij_case234(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t = np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) - t1 = np.arctan(z_bar_k * (1.0 - np.cos(phi_bar_j)) / (np.sin(phi_bar_j) * t)) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) * np.cos(phi_bar_j) - t2 = np.arctanh(z_bar_k / t) - t2_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) * np.sin(phi_bar_j) +def Hphi_phij_case234(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t = xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) + t1 = xp.atan(z_bar_k * (1.0 - xp.cos(phi_bar_j)) / (xp.sin(phi_bar_j) * t)) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.cos(phi_bar_j) + t2 = xp.atanh(z_bar_k / t) + t2_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.sin(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hphi_zk_case234(r, phi_bar_j, theta_M, z_bar_k): - t1 = np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) - t1_coef = np.cos(theta_M) / r - t2 = np.arctanh(r * (1.0 - np.cos(phi_bar_j)) / t1) - t2_coef = np.cos(theta_M) * np.cos(phi_bar_j) +def Hphi_zk_case234(xp, r, phi_bar_j, theta_M, z_bar_k): + t1 = xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) + t1_coef = xp.cos(theta_M) / r + t2 = xp.atanh(r * (1.0 - xp.cos(phi_bar_j)) / t1) + t2_coef = xp.cos(theta_M) * xp.cos(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case234(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) - t1_coef = np.sin(theta_M) * np.sin(phi_bar_M) / r +def Hz_ri_case234(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) / r E = ellipeinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) - E_coef = np.sin(theta_M) * np.cos(phi_bar_M) * np.abs(z_bar_k) / r + E_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.abs(z_bar_k) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r**2 / z_bar_k**2) F_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * (2.0 * r**2 + z_bar_k**2) - / (r * np.abs(z_bar_k)) + / (r * xp.abs(z_bar_k)) ) return t1_coef * t1 + E_coef * E + F_coef * F -def Hz_phij_case234(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.arctanh( +def Hz_phij_case234(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.atanh( r - * (1.0 - np.cos(phi_bar_j)) - / np.sqrt(2.0 * r**2 * (1.0 - np.cos(phi_bar_j)) + z_bar_k**2) + * (1.0 - xp.cos(phi_bar_j)) + / xp.sqrt(2.0 * r**2 * (1.0 - xp.cos(phi_bar_j)) + z_bar_k**2) ) - t1_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) return t1_coef * t1 -def Hz_zk_case234(r, phi_bar_j, theta_M, z_bar_k): - t = np.sqrt(r**2 + z_bar_k**2) +def Hz_zk_case234(xp, r, phi_bar_j, theta_M, z_bar_k): + t = xp.sqrt(r**2 + z_bar_k**2) def Pi(sign): return el3_angle( phi_bar_j / 2.0, 2.0 * r / (r + sign * t), -4.0 * r**2 / z_bar_k**2 ) - Pi_coef = np.cos(theta_M) * np.sign(z_bar_k) + Pi_coef = xp.cos(theta_M) * xp.sign(z_bar_k) return Pi_coef * Pi(1) + Pi_coef * Pi(-1) # 235 ############## -def Hr_ri_case235(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_M) * z_bar_k / (2.0 * r**2) - t2 = np.arctanh( +def Hr_ri_case235(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_M) * z_bar_k / (2.0 * r**2) + t2 = xp.atanh( z_bar_k - / np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) + / xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) ) - t2_coef = np.sin(theta_M) * np.sin(phi_bar_M) / 2.0 * (1.0 - r_i**2 / r**2) + t2_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) / 2.0 * (1.0 - r_i**2 / r**2) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) E_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k - * np.sqrt(r_bar_i**2 + z_bar_k**2) + * xp.sqrt(r_bar_i**2 + z_bar_k**2) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) F_coef = ( - np.sin(theta_M) - * np.cos(phi_bar_M) + xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k * (2.0 * r_i**2 + z_bar_k**2) - / (2.0 * r**2 * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) Pi = el3_angle( phi_bar_j / 2.0, @@ -1605,55 +1609,55 @@ def Hr_ri_case235(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2), ) Pi_coef = ( - np.sin(theta_M) - * np.cos(phi_bar_M) + xp.sin(theta_M) + * xp.cos(phi_bar_M) * z_bar_k * (r**2 + r_i**2) * (r + r_i) - / (2.0 * r**2 * r_bar_i * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * r_bar_i * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) return t1_coef * t1 + t2_coef * t2 + E_coef * E + F_coef * F + Pi_coef * Pi -def Hr_phij_case235(r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.arctanh( +def Hr_phij_case235(xp, r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.atanh( z_bar_k - / np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) + / xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) ) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) * np.cos(phi_bar_j) - t2 = np.arctan( + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.cos(phi_bar_j) + t2 = xp.atan( z_bar_k - * (r * np.cos(phi_bar_j) - r_i) + * (r * xp.cos(phi_bar_j) - r_i) / ( r - * np.sin(phi_bar_j) - * np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) + * xp.sin(phi_bar_j) + * xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) ) ) - t2_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) * np.sin(phi_bar_j) + t2_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.sin(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hr_zk_case235(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): - t1 = np.sin(phi_bar_j) - t1_coef = -np.cos(theta_M) - t2 = np.log( +def Hr_zk_case235(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): + t1 = xp.sin(phi_bar_j) + t1_coef = -xp.cos(theta_M) + t2 = xp.log( r_i - - r * np.cos(phi_bar_j) - + np.sqrt(r_i**2 + r**2 - 2.0 * r_i * r * np.cos(phi_bar_j) + z_bar_k**2) + - r * xp.cos(phi_bar_j) + + xp.sqrt(r_i**2 + r**2 - 2.0 * r_i * r * xp.cos(phi_bar_j) + z_bar_k**2) ) - t2_coef = np.cos(theta_M) * np.sin(phi_bar_j) - t3 = np.arctan(r * np.sin(phi_bar_j) / z_bar_k) - t3_coef = np.cos(theta_M) * z_bar_k / r + t2_coef = xp.cos(theta_M) * xp.sin(phi_bar_j) + t3 = xp.atan(r * xp.sin(phi_bar_j) / z_bar_k) + t3_coef = xp.cos(theta_M) * z_bar_k / r E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) - E_coef = np.cos(theta_M) * np.sqrt(r_bar_i**2 + z_bar_k**2) / r + E_coef = xp.cos(theta_M) * xp.sqrt(r_bar_i**2 + z_bar_k**2) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) F_coef = ( - -np.cos(theta_M) + -xp.cos(theta_M) * (r**2 + r_i**2 + z_bar_k**2) - / (r * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (r * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) - t = np.sqrt(r**2 + z_bar_k**2) + t = xp.sqrt(r**2 + z_bar_k**2) def Pi1(sign): return el3_angle( @@ -1664,8 +1668,8 @@ def Pi1(sign): def Pi1_coef(sign): return ( - -np.cos(theta_M) - / (r * np.sqrt((r**2 + z_bar_k**2) * (r_bar_i**2 + z_bar_k**2))) + -xp.cos(theta_M) + / (r * xp.sqrt((r**2 + z_bar_k**2) * (r_bar_i**2 + z_bar_k**2))) * (t - sign * r) * (r_i + sign * t) ** 2 ) @@ -1673,7 +1677,7 @@ def Pi1_coef(sign): def Pi2(sign): return el3_angle( arctan_k_tan_2( - np.sqrt(((r_i + r) ** 2 + z_bar_k**2) / (r_bar_i**2 + z_bar_k**2)), + xp.sqrt(((r_i + r) ** 2 + z_bar_k**2) / (r_bar_i**2 + z_bar_k**2)), phi_bar_j, ), 1.0 @@ -1686,12 +1690,12 @@ def Pi2(sign): def Pi2_coef(sign): return ( sign - * np.cos(theta_M) + * xp.cos(theta_M) * z_bar_k**2 * (r_bar_i**2 + z_bar_k**2) / ( r - * np.sqrt((r**2 + z_bar_k**2) * ((r + r_i) ** 2 + z_bar_k**2)) + * xp.sqrt((r**2 + z_bar_k**2) * ((r + r_i) ** 2 + z_bar_k**2)) * (r + sign * t) ) ) @@ -1709,29 +1713,29 @@ def Pi2_coef(sign): ) -def Hphi_ri_case235(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - t1 = np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) - t1_coef = -np.sin(theta_M) * np.cos(phi_bar_M) * z_bar_k / (2.0 * r**2) - t2 = np.arctanh( +def Hphi_ri_case235(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + t1 = xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) + t1_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) * z_bar_k / (2.0 * r**2) + t2 = xp.atanh( z_bar_k - / np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) + / xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) ) - t2_coef = -np.sin(theta_M) * np.cos(phi_bar_M) * (r**2 + r_i**2) / (2.0 * r**2) + t2_coef = -xp.sin(theta_M) * xp.cos(phi_bar_M) * (r**2 + r_i**2) / (2.0 * r**2) E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) E_coef = ( - np.sin(theta_M) - * np.sin(phi_bar_M) + xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k - * np.sqrt(r_bar_i**2 + z_bar_k**2) + * xp.sqrt(r_bar_i**2 + z_bar_k**2) / (2.0 * r**2) ) F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2)) F_coef = ( - -np.sin(theta_M) - * np.sin(phi_bar_M) + -xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k * (2.0 * r**2 + 2.0 * r_i**2 + z_bar_k**2) - / (2.0 * r**2 * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) Pi = el3_angle( phi_bar_j / 2.0, @@ -1739,61 +1743,59 @@ def Hphi_ri_case235(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): -4.0 * r * r_i / (r_bar_i**2 + z_bar_k**2), ) Pi_coef = ( - np.sin(theta_M) - * np.sin(phi_bar_M) + xp.sin(theta_M) + * xp.sin(phi_bar_M) * z_bar_k * (r + r_i) ** 2 - / (2.0 * r**2 * np.sqrt(r_bar_i**2 + z_bar_k**2)) + / (2.0 * r**2 * xp.sqrt(r_bar_i**2 + z_bar_k**2)) ) return t1_coef * t1 + t2_coef * t2 + E_coef * E + F_coef * F + Pi_coef * Pi -def Hphi_phij_case235(r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t = np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) - t1 = np.arctan( - z_bar_k * (r * np.cos(phi_bar_j) - r_i) / (r * np.sin(phi_bar_j) * t) - ) - t1_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) * np.cos(phi_bar_j) - t2 = np.arctanh(z_bar_k / t) - t2_coef = np.sin(theta_M) * np.sin(phi_bar_Mj) * np.sin(phi_bar_j) +def Hphi_phij_case235(xp, r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t = xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) + t1 = xp.atan(z_bar_k * (r * xp.cos(phi_bar_j) - r_i) / (r * xp.sin(phi_bar_j) * t)) + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.cos(phi_bar_j) + t2 = xp.atanh(z_bar_k / t) + t2_coef = xp.sin(theta_M) * xp.sin(phi_bar_Mj) * xp.sin(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hphi_zk_case235(r, r_i, phi_bar_j, theta_M, z_bar_k): - t1 = np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) - t1_coef = np.cos(theta_M) / r - t2 = np.arctanh((r * np.cos(phi_bar_j) - r_i) / t1) - t2_coef = -np.cos(theta_M) * np.cos(phi_bar_j) +def Hphi_zk_case235(xp, r, r_i, phi_bar_j, theta_M, z_bar_k): + t1 = xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) + t1_coef = xp.cos(theta_M) / r + t2 = xp.atanh((r * xp.cos(phi_bar_j) - r_i) / t1) + t2_coef = -xp.cos(theta_M) * xp.cos(phi_bar_j) return t1_coef * t1 + t2_coef * t2 -def Hz_ri_case235(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): +def Hz_ri_case235(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): t = r_bar_i**2 + z_bar_k**2 - t1 = np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) - t1_coef = np.sin(theta_M) * np.sin(phi_bar_M) / r + t1 = xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) + t1_coef = xp.sin(theta_M) * xp.sin(phi_bar_M) / r E = ellipeinc(phi_bar_j / 2.0, -4.0 * r * r_i / t) - E_coef = np.sin(theta_M) * np.cos(phi_bar_M) * np.sqrt(t) / r + E_coef = xp.sin(theta_M) * xp.cos(phi_bar_M) * xp.sqrt(t) / r F = ellipkinc(phi_bar_j / 2.0, -4.0 * r * r_i / t) F_coef = ( - -np.sin(theta_M) - * np.cos(phi_bar_M) + -xp.sin(theta_M) + * xp.cos(phi_bar_M) * (r**2 + r_i**2 + z_bar_k**2) - / (r * np.sqrt(t)) + / (r * xp.sqrt(t)) ) return t1_coef * t1 + E_coef * E + F_coef * F -def Hz_phij_case235(r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - t1 = np.arctanh( - (r * np.cos(phi_bar_j) - r_i) - / np.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * np.cos(phi_bar_j) + z_bar_k**2) +def Hz_phij_case235(xp, r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + t1 = xp.atanh( + (r * xp.cos(phi_bar_j) - r_i) + / xp.sqrt(r**2 + r_i**2 - 2.0 * r * r_i * xp.cos(phi_bar_j) + z_bar_k**2) ) - t1_coef = -np.sin(theta_M) * np.sin(phi_bar_Mj) + t1_coef = -xp.sin(theta_M) * xp.sin(phi_bar_Mj) return t1_coef * t1 -def Hz_zk_case235(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): - t = np.sqrt(r**2 + z_bar_k**2) +def Hz_zk_case235(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k): + t = xp.sqrt(r**2 + z_bar_k**2) def Pi(sign): return el3_angle( @@ -1804,10 +1806,10 @@ def Pi(sign): def Pi_coef(sign): return ( - np.cos(theta_M) + xp.cos(theta_M) * z_bar_k * (r_i + sign * t) - / (np.sqrt(r_bar_i**2 + z_bar_k**2) * (r + sign * t)) + / (xp.sqrt(r_bar_i**2 + z_bar_k**2) * (r + sign * t)) ) return Pi_coef(1) * Pi(1) + Pi_coef(-1) * Pi(-1) @@ -1823,286 +1825,294 @@ def Pi_coef(sign): # out: ndarray, shape (n,3,3) # (n)vector, (3)r_phi_z, (3)face -def case112(r_i, phi_bar_M, theta_M): - results = np.zeros((len(r_i), 3, 3)) - results[:, 1, 2] = Hphi_zk_case112(r_i, theta_M) - results[:, 2, 0] = Hz_ri_case112(phi_bar_M, theta_M) - results[:, 2, 1] = Hz_phij_case112(r_i, phi_bar_M, theta_M) +def case112(xp, r_i, phi_bar_M, theta_M): + results = xp.zeros(((r_i.shape[0]), 3, 3)) + results[:, 1, 2] = Hphi_zk_case112(xp, r_i, theta_M) + results[:, 2, 0] = Hz_ri_case112(xp, phi_bar_M, theta_M) + results[:, 2, 1] = Hz_phij_case112(xp, r_i, phi_bar_M, theta_M) return results -def case113(r, phi_bar_M, theta_M): - results = np.zeros((len(r), 3, 3)) - results[:, 1, 2] = Hphi_zk_case113(r, theta_M) - results[:, 2, 1] = Hz_phij_case113(r, phi_bar_M, theta_M) +def case113(xp, r, phi_bar_M, theta_M): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 1, 2] = Hphi_zk_case113(xp, r, theta_M) + results[:, 2, 1] = Hz_phij_case113(xp, r, phi_bar_M, theta_M) return results -def case115(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 2] = Hr_zk_case115(r, r_i, r_bar_i, phi_bar_j, theta_M) - results[:, 1, 2] = Hphi_zk_case115(r, r_i, r_bar_i, theta_M) - results[:, 2, 0] = Hz_ri_case115(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M) - results[:, 2, 1] = Hz_phij_case115(r_bar_i, phi_bar_M, theta_M) +def case115(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 2] = Hr_zk_case115(xp, r, r_i, r_bar_i, phi_bar_j, theta_M) + results[:, 1, 2] = Hphi_zk_case115(xp, r, r_i, r_bar_i, theta_M) + results[:, 2, 0] = Hz_ri_case115(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M) + results[:, 2, 1] = Hz_phij_case115(xp, r_bar_i, phi_bar_M, theta_M) return results -def case122(r_i, phi_bar_M, theta_M): - results = np.zeros((len(r_i), 3, 3)) - results[:, 1, 2] = Hphi_zk_case122(r_i, theta_M) - results[:, 2, 0] = Hz_ri_case122(phi_bar_M, theta_M) - results[:, 2, 1] = Hz_phij_case122(r_i, phi_bar_M, theta_M) +def case122(xp, r_i, phi_bar_M, theta_M): + results = xp.zeros(((r_i.shape[0]), 3, 3)) + results[:, 1, 2] = Hphi_zk_case122(xp, r_i, theta_M) + results[:, 2, 0] = Hz_ri_case122(xp, phi_bar_M, theta_M) + results[:, 2, 1] = Hz_phij_case122(xp, r_i, phi_bar_M, theta_M) return results -def case123(r, phi_bar_M, theta_M): - results = np.zeros((len(r), 3, 3)) - results[:, 1, 2] = Hphi_zk_case123(r, theta_M) - results[:, 2, 1] = Hz_phij_case123(r, phi_bar_M, theta_M) +def case123(xp, r, phi_bar_M, theta_M): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 1, 2] = Hphi_zk_case123(xp, r, theta_M) + results[:, 2, 1] = Hz_phij_case123(xp, r, phi_bar_M, theta_M) return results -def case124(r, phi_bar_M, theta_M): - results = np.zeros((len(r), 3, 3)) - results[:, 1, 2] = Hphi_zk_case124(r, theta_M) - results[:, 2, 0] = Hz_ri_case124(phi_bar_M, theta_M) - results[:, 2, 1] = Hz_phij_case124(r, phi_bar_M, theta_M) +def case124(xp, r, phi_bar_M, theta_M): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 1, 2] = Hphi_zk_case124(xp, r, theta_M) + results[:, 2, 0] = Hz_ri_case124(xp, phi_bar_M, theta_M) + results[:, 2, 1] = Hz_phij_case124(xp, r, phi_bar_M, theta_M) return results -def case125(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 2] = Hr_zk_case125(r, r_i, r_bar_i, phi_bar_j, theta_M) - results[:, 1, 2] = Hphi_zk_case125(r, r_i, theta_M) - results[:, 2, 0] = Hz_ri_case125(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M) - results[:, 2, 1] = Hz_phij_case125(r, r_i, phi_bar_M, theta_M) +def case125(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 2] = Hr_zk_case125(xp, r, r_i, r_bar_i, phi_bar_j, theta_M) + results[:, 1, 2] = Hphi_zk_case125(xp, r, r_i, theta_M) + results[:, 2, 0] = Hz_ri_case125(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M) + results[:, 2, 1] = Hz_phij_case125(xp, r, r_i, phi_bar_M, theta_M) return results -def case132(r, r_i, phi_bar_j, phi_bar_Mj, theta_M): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 2] = Hr_zk_case132(r_i, phi_bar_j, theta_M) - results[:, 1, 2] = Hphi_zk_case132(r_i, phi_bar_j, theta_M) - results[:, 2, 0] = Hz_ri_case132(phi_bar_Mj, theta_M) - results[:, 2, 1] = Hz_phij_case132(r_i, phi_bar_Mj, theta_M) +def case132(xp, r, r_i, phi_bar_j, phi_bar_Mj, theta_M): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 2] = Hr_zk_case132(xp, r_i, phi_bar_j, theta_M) + results[:, 1, 2] = Hphi_zk_case132(xp, r_i, phi_bar_j, theta_M) + results[:, 2, 0] = Hz_ri_case132(xp, phi_bar_Mj, theta_M) + results[:, 2, 1] = Hz_phij_case132(xp, r_i, phi_bar_Mj, theta_M) return results -def case133(r, phi_bar_j, phi_bar_Mj, theta_M): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 2] = Hr_zk_case133(r, phi_bar_j, theta_M) - results[:, 1, 2] = Hphi_zk_case133(phi_bar_j, theta_M) - results[:, 2, 1] = Hz_phij_case133(phi_bar_j, phi_bar_Mj, theta_M) +def case133(xp, r, phi_bar_j, phi_bar_Mj, theta_M): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 2] = Hr_zk_case133(xp, r, phi_bar_j, theta_M) + results[:, 1, 2] = Hphi_zk_case133(xp, phi_bar_j, theta_M) + results[:, 2, 1] = Hz_phij_case133(xp, phi_bar_j, phi_bar_Mj, theta_M) return results -def case134(r, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 2] = Hr_zk_case134(r, phi_bar_j, theta_M) - results[:, 1, 2] = Hphi_zk_case134(phi_bar_j, theta_M) - results[:, 2, 0] = Hz_ri_case134(phi_bar_j, phi_bar_M, theta_M) - results[:, 2, 1] = Hz_phij_case134(phi_bar_j, phi_bar_Mj, theta_M) +def case134(xp, r, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 2] = Hr_zk_case134(xp, r, phi_bar_j, theta_M) + results[:, 1, 2] = Hphi_zk_case134(xp, phi_bar_j, theta_M) + results[:, 2, 0] = Hz_ri_case134(xp, phi_bar_j, phi_bar_M, theta_M) + results[:, 2, 1] = Hz_phij_case134(xp, phi_bar_j, phi_bar_Mj, theta_M) return results -def case135(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 2] = Hr_zk_case135(r, r_i, r_bar_i, phi_bar_j, theta_M) - results[:, 1, 2] = Hphi_zk_case135(r, r_i, phi_bar_j, theta_M) - results[:, 2, 0] = Hz_ri_case135(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M) - results[:, 2, 1] = Hz_phij_case135(r, r_i, phi_bar_j, phi_bar_Mj, theta_M) +def case135(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 2] = Hr_zk_case135(xp, r, r_i, r_bar_i, phi_bar_j, theta_M) + results[:, 1, 2] = Hphi_zk_case135(xp, r, r_i, phi_bar_j, theta_M) + results[:, 2, 0] = Hz_ri_case135(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M) + results[:, 2, 1] = Hz_phij_case135(xp, r, r_i, phi_bar_j, phi_bar_Mj, theta_M) return results -def case211(phi_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(phi_j), 3, 3)) - results[:, 0, 1] = Hr_phij_case211(phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case211(phi_j, theta_M, z_bar_k) +def case211(xp, phi_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((phi_j.shape[0]), 3, 3)) + results[:, 0, 1] = Hr_phij_case211(xp, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case211(xp, phi_j, theta_M, z_bar_k) return results -def case212(r_i, phi_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(r_i), 3, 3)) - results[:, 0, 0] = Hr_ri_case212(r_i, phi_j, phi_bar_M, theta_M, z_bar_k) - results[:, 0, 1] = Hr_phij_case212(r_i, phi_bar_M, theta_M, z_bar_k) - results[:, 1, 0] = Hphi_ri_case212(r_i, phi_j, phi_bar_M, theta_M, z_bar_k) - results[:, 1, 2] = Hphi_zk_case212(r_i, theta_M, z_bar_k) - results[:, 2, 0] = Hz_ri_case212(r_i, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 1] = Hz_phij_case212(r_i, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case212(r_i, phi_j, theta_M, z_bar_k) +def case212(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((r_i.shape[0]), 3, 3)) + results[:, 0, 0] = Hr_ri_case212(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k) + results[:, 0, 1] = Hr_phij_case212(xp, r_i, phi_bar_M, theta_M, z_bar_k) + results[:, 1, 0] = Hphi_ri_case212(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case212(xp, r_i, theta_M, z_bar_k) + results[:, 2, 0] = Hz_ri_case212(xp, r_i, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case212(xp, r_i, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case212(xp, r_i, phi_j, theta_M, z_bar_k) return results -def case213(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 1] = Hr_phij_case213(r, phi_bar_M, theta_M, z_bar_k) - results[:, 1, 2] = Hphi_zk_case213(r, theta_M, z_bar_k) - results[:, 2, 1] = Hz_phij_case213(r, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case213(phi_bar_j, theta_M, z_bar_k) +def case213(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 1] = Hr_phij_case213(xp, r, phi_bar_M, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case213(xp, r, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case213(xp, r, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case213(xp, phi_bar_j, theta_M, z_bar_k) return results -def case214(r, phi_j, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 0] = Hr_ri_case214(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) - results[:, 0, 1] = Hr_phij_case214(phi_bar_M, theta_M, z_bar_k) - results[:, 0, 2] = Hr_zk_case214(r, phi_bar_j, theta_M, z_bar_k) - results[:, 1, 0] = Hphi_ri_case214(r, phi_j, phi_bar_j, phi_bar_M, theta_M, z_bar_k) - results[:, 1, 2] = Hphi_zk_case214(r, theta_M, z_bar_k) - results[:, 2, 0] = Hz_ri_case214(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case214(r, phi_bar_j, theta_M, z_bar_k) +def case214(xp, r, phi_j, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 0] = Hr_ri_case214(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) + results[:, 0, 1] = Hr_phij_case214(xp, phi_bar_M, theta_M, z_bar_k) + results[:, 0, 2] = Hr_zk_case214(xp, r, phi_bar_j, theta_M, z_bar_k) + results[:, 1, 0] = Hphi_ri_case214( + xp, r, phi_j, phi_bar_j, phi_bar_M, theta_M, z_bar_k + ) + results[:, 1, 2] = Hphi_zk_case214(xp, r, theta_M, z_bar_k) + results[:, 2, 0] = Hz_ri_case214(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case214(xp, r, phi_bar_j, theta_M, z_bar_k) return results -def case215(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(r), 3, 3)) +def case215(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((r.shape[0]), 3, 3)) results[:, 0, 0] = Hr_ri_case215( - r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k ) - results[:, 0, 1] = Hr_phij_case215(r_bar_i, phi_bar_M, theta_M, z_bar_k) - results[:, 0, 2] = Hr_zk_case215(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) + results[:, 0, 1] = Hr_phij_case215(xp, r_bar_i, phi_bar_M, theta_M, z_bar_k) + results[:, 0, 2] = Hr_zk_case215(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) results[:, 1, 0] = Hphi_ri_case215( - r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k ) - results[:, 1, 2] = Hphi_zk_case215(r, r_bar_i, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case215(xp, r, r_bar_i, theta_M, z_bar_k) results[:, 2, 0] = Hz_ri_case215( - r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k ) - results[:, 2, 1] = Hz_phij_case215(r_bar_i, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case215(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case215(xp, r_bar_i, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case215(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) return results -def case221(phi_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(phi_j), 3, 3)) - results[:, 0, 1] = Hr_phij_case221(phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case221(phi_j, theta_M, z_bar_k) +def case221(xp, phi_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((phi_j.shape[0]), 3, 3)) + results[:, 0, 1] = Hr_phij_case221(xp, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case221(xp, phi_j, theta_M, z_bar_k) return results -def case222(r_i, phi_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(r_i), 3, 3)) - results[:, 0, 0] = Hr_ri_case222(r_i, phi_j, phi_bar_M, theta_M, z_bar_k) - results[:, 0, 1] = Hr_phij_case222(r_i, phi_bar_M, theta_M, z_bar_k) - results[:, 1, 0] = Hphi_ri_case222(r_i, phi_j, phi_bar_M, theta_M, z_bar_k) - results[:, 1, 2] = Hphi_zk_case222(r_i, theta_M, z_bar_k) - results[:, 2, 0] = Hz_ri_case222(r_i, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 1] = Hz_phij_case222(r_i, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case222(r_i, phi_j, theta_M, z_bar_k) +def case222(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((r_i.shape[0]), 3, 3)) + results[:, 0, 0] = Hr_ri_case222(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k) + results[:, 0, 1] = Hr_phij_case222(xp, r_i, phi_bar_M, theta_M, z_bar_k) + results[:, 1, 0] = Hphi_ri_case222(xp, r_i, phi_j, phi_bar_M, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case222(xp, r_i, theta_M, z_bar_k) + results[:, 2, 0] = Hz_ri_case222(xp, r_i, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case222(xp, r_i, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case222(xp, r_i, phi_j, theta_M, z_bar_k) return results -def case223(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 1] = Hr_phij_case223(r, phi_bar_M, theta_M, z_bar_k) - results[:, 1, 2] = Hphi_zk_case223(r, theta_M, z_bar_k) - results[:, 2, 1] = Hz_phij_case223(r, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case223(r, phi_bar_j, theta_M, z_bar_k) +def case223(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 1] = Hr_phij_case223(xp, r, phi_bar_M, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case223(xp, r, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case223(xp, r, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case223(xp, r, phi_bar_j, theta_M, z_bar_k) return results -def case224(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 0] = Hr_ri_case224(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) - results[:, 0, 1] = Hr_phij_case224(r, phi_bar_M, theta_M, z_bar_k) - results[:, 0, 2] = Hr_zk_case224(r, phi_bar_j, theta_M, z_bar_k) - results[:, 1, 0] = Hphi_ri_case224(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) - results[:, 1, 2] = Hphi_zk_case224(r, theta_M, z_bar_k) - results[:, 2, 0] = Hz_ri_case224(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 1] = Hz_phij_case224(r, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case224(r, phi_bar_j, theta_M, z_bar_k) +def case224(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 0] = Hr_ri_case224(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) + results[:, 0, 1] = Hr_phij_case224(xp, r, phi_bar_M, theta_M, z_bar_k) + results[:, 0, 2] = Hr_zk_case224(xp, r, phi_bar_j, theta_M, z_bar_k) + results[:, 1, 0] = Hphi_ri_case224(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case224(xp, r, theta_M, z_bar_k) + results[:, 2, 0] = Hz_ri_case224(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case224(xp, r, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case224(xp, r, phi_bar_j, theta_M, z_bar_k) return results -def case225(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): - results = np.zeros((len(r), 3, 3)) +def case225(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k): + results = xp.zeros(((r.shape[0]), 3, 3)) results[:, 0, 0] = Hr_ri_case225( - r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k ) - results[:, 0, 1] = Hr_phij_case225(r, r_i, phi_bar_M, theta_M, z_bar_k) - results[:, 0, 2] = Hr_zk_case225(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) + results[:, 0, 1] = Hr_phij_case225(xp, r, r_i, phi_bar_M, theta_M, z_bar_k) + results[:, 0, 2] = Hr_zk_case225(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) results[:, 1, 0] = Hphi_ri_case225( - r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k ) - results[:, 1, 2] = Hphi_zk_case225(r, r_i, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case225(xp, r, r_i, theta_M, z_bar_k) results[:, 2, 0] = Hz_ri_case225( - r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k ) - results[:, 2, 1] = Hz_phij_case225(r, r_i, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case225(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case225(xp, r, r_i, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case225(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) return results -def case231(phi_j, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - results = np.zeros((len(phi_j), 3, 3)) - results[:, 0, 1] = Hr_phij_case231(phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 1, 1] = Hphi_phij_case231(phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case231(phi_j, theta_M, z_bar_k) +def case231(xp, phi_j, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + results = xp.zeros(((phi_j.shape[0]), 3, 3)) + results[:, 0, 1] = Hr_phij_case231(xp, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) + results[:, 1, 1] = Hphi_phij_case231(xp, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case231(xp, phi_j, theta_M, z_bar_k) return results -def case232(r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): - results = np.zeros((len(r_i), 3, 3)) +def case232(xp, r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): + results = xp.zeros(((r_i.shape[0]), 3, 3)) results[:, 0, 0] = Hr_ri_case232( - r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k + xp, r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k ) - results[:, 0, 1] = Hr_phij_case232(r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 0, 2] = Hr_zk_case232(r_i, phi_bar_j, theta_M, z_bar_k) + results[:, 0, 1] = Hr_phij_case232(xp, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) + results[:, 0, 2] = Hr_zk_case232(xp, r_i, phi_bar_j, theta_M, z_bar_k) results[:, 1, 0] = Hphi_ri_case232( - r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k + xp, r_i, phi_j, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k + ) + results[:, 1, 1] = Hphi_phij_case232( + xp, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k ) - results[:, 1, 1] = Hphi_phij_case232(r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 1, 2] = Hphi_zk_case232(r_i, phi_bar_j, theta_M, z_bar_k) - results[:, 2, 0] = Hz_ri_case232(r_i, phi_bar_Mj, theta_M, z_bar_k) - results[:, 2, 1] = Hz_phij_case232(r_i, phi_bar_Mj, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case232(r_i, phi_j, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case232(xp, r_i, phi_bar_j, theta_M, z_bar_k) + results[:, 2, 0] = Hz_ri_case232(xp, r_i, phi_bar_Mj, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case232(xp, r_i, phi_bar_Mj, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case232(xp, r_i, phi_j, theta_M, z_bar_k) return results -def case233(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 1] = Hr_phij_case233(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 0, 2] = Hr_zk_case233(r, phi_bar_j, theta_M, z_bar_k) - results[:, 1, 1] = Hphi_phij_case233(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 1, 2] = Hphi_zk_case233(r, phi_bar_j, theta_M, z_bar_k) - results[:, 2, 1] = Hz_phij_case233(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case233(r, phi_bar_j, theta_M, z_bar_k) +def case233(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 1] = Hr_phij_case233(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) + results[:, 0, 2] = Hr_zk_case233(xp, r, phi_bar_j, theta_M, z_bar_k) + results[:, 1, 1] = Hphi_phij_case233(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case233(xp, r, phi_bar_j, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case233(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case233(xp, r, phi_bar_j, theta_M, z_bar_k) return results -def case234(r, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): - results = np.zeros((len(r), 3, 3)) - results[:, 0, 0] = Hr_ri_case234(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) - results[:, 0, 1] = Hr_phij_case234(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 0, 2] = Hr_zk_case234(r, phi_bar_j, theta_M, z_bar_k) - results[:, 1, 0] = Hphi_ri_case234(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) - results[:, 1, 1] = Hphi_phij_case234(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 1, 2] = Hphi_zk_case234(r, phi_bar_j, theta_M, z_bar_k) - results[:, 2, 0] = Hz_ri_case234(r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) - results[:, 2, 1] = Hz_phij_case234(r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case234(r, phi_bar_j, theta_M, z_bar_k) +def case234(xp, r, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): + results = xp.zeros(((r.shape[0]), 3, 3)) + results[:, 0, 0] = Hr_ri_case234(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) + results[:, 0, 1] = Hr_phij_case234(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) + results[:, 0, 2] = Hr_zk_case234(xp, r, phi_bar_j, theta_M, z_bar_k) + results[:, 1, 0] = Hphi_ri_case234(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) + results[:, 1, 1] = Hphi_phij_case234(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case234(xp, r, phi_bar_j, theta_M, z_bar_k) + results[:, 2, 0] = Hz_ri_case234(xp, r, phi_bar_j, phi_bar_M, theta_M, z_bar_k) + results[:, 2, 1] = Hz_phij_case234(xp, r, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case234(xp, r, phi_bar_j, theta_M, z_bar_k) return results -def case235(r, r_i, r_bar_i, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): - results = np.zeros((len(r), 3, 3)) +def case235(xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, phi_bar_Mj, theta_M, z_bar_k): + results = xp.zeros(((r.shape[0]), 3, 3)) results[:, 0, 0] = Hr_ri_case235( - r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k ) - results[:, 0, 1] = Hr_phij_case235(r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 0, 2] = Hr_zk_case235(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) + results[:, 0, 1] = Hr_phij_case235( + xp, r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k + ) + results[:, 0, 2] = Hr_zk_case235(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) results[:, 1, 0] = Hphi_ri_case235( - r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k ) results[:, 1, 1] = Hphi_phij_case235( - r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k + xp, r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k ) - results[:, 1, 2] = Hphi_zk_case235(r, r_i, phi_bar_j, theta_M, z_bar_k) + results[:, 1, 2] = Hphi_zk_case235(xp, r, r_i, phi_bar_j, theta_M, z_bar_k) results[:, 2, 0] = Hz_ri_case235( - r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + xp, r, r_i, r_bar_i, phi_bar_j, phi_bar_M, theta_M, z_bar_k + ) + results[:, 2, 1] = Hz_phij_case235( + xp, r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k ) - results[:, 2, 1] = Hz_phij_case235(r, r_i, phi_bar_j, phi_bar_Mj, theta_M, z_bar_k) - results[:, 2, 2] = Hz_zk_case235(r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) + results[:, 2, 2] = Hz_zk_case235(xp, r, r_i, r_bar_i, phi_bar_j, theta_M, z_bar_k) return results @@ -2146,11 +2156,11 @@ def magnet_cylinder_segment_Hfield( >>> import numpy as np >>> import magpylib as magpy >>> B = magpy.core.magnet_cylinder_segment_Hfield( - ... observers=np.array([(1,1,2), (0,0,0)]), - ... dimensions=np.array([(1,2,.1,.2,-1,1), (1,2,.3,.9,0,1)]), - ... magnetizations=np.array([(1e7,.1,.2), (1e6,1.1,2.2)]), + ... observers=xp.array([(1,1,2), (0,0,0)]), + ... dimensions=xp.array([(1,2,.1,.2,-1,1), (1,2,.3,.9,0,1)]), + ... magnetizations=xp.array([(1e7,.1,.2), (1e6,1.1,2.2)]), ... ) - >>> with np.printoptions(precision=3): + >>> with xp.printoptions(precision=3): ... print(B) [[-1948.144 32319.944 17616.886] [14167.65 1419.941 17921.646]] @@ -2161,22 +2171,28 @@ def magnet_cylinder_segment_Hfield( Materials, Volume 559, 1 October 2022, 169482 """ + xp = array_namespace(observers, dimensions, magnetizations) + observers, dimensions, magnetizations = xp_promote( + observers, dimensions, magnetizations, force_floating=True, xp=xp + ) # tile inputs into 8-stacks (boundary cases) - r, phi, z = np.repeat(observers, 8, axis=0).T - r_i = np.repeat(dimensions[:, :2], 4) - phi_j = np.repeat(np.tile(dimensions[:, 2:4], 2), 2) - z_k = np.ravel(np.tile(dimensions[:, 4:6], 4)) - _, phi_M, theta_M = np.repeat(magnetizations, 8, axis=0).T + rphiz = xp.repeat(observers, 8, axis=0).T + r, phi, z = (rphiz[i, ...] for i in range(3)) + r_i = xp.repeat(dimensions[:, :2], 4) + phi_j = xp.repeat(xp.tile(dimensions[:, 2:4], (2,)), 2) + z_k = xp.reshape(xp.tile(dimensions[:, 4:6], (4,)), (-1,)) + rphiz_M = xp.repeat(magnetizations, 8, axis=0).T + _, phi_M, theta_M = (rphiz_M[i, ...] for i in range(3)) # initialize results array with nan - result = np.empty((len(r), 3, 3)) - result[:] = np.nan + result = xp.empty(((r.shape[0]), 3, 3)) + result[:, ...] = xp.nan # cases to evaluate cases = determine_cases(r, phi, z, r_i, phi_j, z_k) # list of all possible cases - excluding the nan-cases 111, 114, 121, 131 - case_id = np.array( + case_id = xp.asarray( [ 112, 113, @@ -2288,11 +2304,15 @@ def magnet_cylinder_segment_Hfield( for cid, cfkt, cargs in zip(case_id, case_fkt, case_args, strict=False): mask = cases == cid if any(mask): - result[mask] = cfkt(*[allargs[aid][mask] for aid in cargs]) + result[mask] = cfkt(xp, *[allargs[aid][mask] for aid in cargs]) # sum up contributions from different boundary cases (ax1) and different face types (ax3) - result = np.reshape(result, (-1, 8, 3, 3)) - result = np.sum(result[:, (1, 2, 4, 7)] - result[:, (0, 3, 5, 6)], axis=(1, 3)) + result = xp.reshape(result, (-1, 8, 3, 3)) + result0 = result[:, 1, ...] - result[:, 0, ...] + result1 = result[:, 2, ...] - result[:, 3, ...] + result2 = result[:, 4, ...] - result[:, 5, ...] + result3 = result[:, 7, ...] - result[:, 6, ...] + result = xp.sum(result0 + result1 + result2 + result3, axis=2) # multiply with magnetization amplitude result = result.T * magnetizations[:, 0] * 1e-7 / MU0 @@ -2312,8 +2332,9 @@ def BHJM_cylinder_segment_internal( Falls back to magnet_cylinder_field whenever the section angles describe the full 360° cylinder. """ + xp = array_namespace(observers, polarization, dimension) - BHfinal = np.zeros_like(observers, dtype=float) + BHfinal = xp.zeros_like(observers, dtype=float) r1, r2, h, phi1, phi2 = dimension.T @@ -2333,7 +2354,7 @@ def BHJM_cylinder_segment_internal( field=field, observers=observers[mask1x], polarization=polarization[mask1x], - dimension=np.c_[2 * r2[mask1x], h[mask1x]], + dimension=xp.c_[2 * r2[mask1x], h[mask1x]], ) # case2a: hollow cylinder <- should be vectorized together with above @@ -2342,7 +2363,7 @@ def BHJM_cylinder_segment_internal( field=field, observers=observers[mask2], polarization=polarization[mask2], - dimension=np.c_[2 * r1[mask2], h[mask2]], + dimension=xp.c_[2 * r1[mask2], h[mask2]], ) return BHfinal From d5dcaad934eb9c76be977988e7f99f628547e0e1 Mon Sep 17 00:00:00 2001 From: purepani Date: Tue, 20 May 2025 16:26:39 -0500 Subject: [PATCH 15/16] Utilities. lazy_while was taken from the lazy_apply implementation in https://github.com/data-apis/array-api-extra, whereas the rest are taken from https://github.com/scipy/scipy/blob/main/scipy/_lib/_array_api.py --- src/magpylib/_src/array_api_utils.py | 254 ++++++++++++++++++++++++++- 1 file changed, 248 insertions(+), 6 deletions(-) diff --git a/src/magpylib/_src/array_api_utils.py b/src/magpylib/_src/array_api_utils.py index 89af5adff..f4e3e4d42 100644 --- a/src/magpylib/_src/array_api_utils.py +++ b/src/magpylib/_src/array_api_utils.py @@ -1,16 +1,34 @@ +from __future__ import annotations + +import math +from collections.abc import Callable, Sequence +from functools import partial, wraps +from types import ModuleType +from typing import Any, Literal, ParamSpec, TypeAlias, cast + +import array_api_compat as _compat import numpy as np import numpy.typing as npt +from array_api_compat import ( + array_namespace, + is_dask_namespace, + is_jax_namespace, + is_numpy_array, +) from array_api_compat import ( is_numpy_namespace as is_numpy, +) +from array_api_compat import ( is_torch_namespace as is_torch, - array_namespace, ) -from typing import Any, TypeAlias, Literal -from types import ModuleType +from typing_extensions import TypeIs Array: TypeAlias = Any # To be changed to a Protocol later (see array-api#589) +DType: TypeAlias = Any # To be changed to a Protocol later (see array-api#589) ArrayLike: TypeAlias = Array | npt.ArrayLike +P = ParamSpec("P") + def _check_finite(array: Array, xp: ModuleType) -> None: """Check for NaNs or Infs.""" @@ -71,9 +89,8 @@ def xp_default_dtype(xp): if is_torch(xp): # historically, we allow pytorch to keep its default of float32 return xp.get_default_dtype() - else: - # we default to float64 - return xp.float64 + # we default to float64 + return xp.float64 def xp_result_type(*args, force_floating=False, xp): @@ -192,3 +209,228 @@ def xp_promote(*args, broadcast=False, force_floating=False, xp): out.append(arg) return out[0] if len(out) == 1 else tuple(out) + + +def _is_jax_jit_enabled(xp: ModuleType) -> bool: # numpydoc ignore=PR01,RT01 + """Return True if this function is being called inside ``jax.jit``.""" + import jax # pylint: disable=import-outside-toplevel + + x = xp.asarray(False) + try: + return bool(x) + except jax.errors.TracerBoolConversionError: + return True + + +def is_python_scalar(x: object) -> TypeIs[complex]: # numpydoc ignore=PR01,RT01 + """Return True if `x` is a Python scalar, False otherwise.""" + # isinstance(x, float) returns True for np.float64 + # isinstance(x, complex) returns True for np.complex128 + # bool is a subclass of int + return isinstance(x, int | float | complex) and not is_numpy_array(x) + + +def lazy_while( # type: ignore[valid-type] # numpydoc ignore=GL07,SA04 + func: Callable[P, Array | ArrayLike | Sequence[Array | ArrayLike]], + cond: Callable[P, bool], + args: Array | complex | None, + shape: tuple[int | None, ...] | Sequence[tuple[int | None, ...]] | None = None, + dtype: DType | Sequence[DType] | None = None, + as_numpy: bool = False, + xp: ModuleType | None = None, + **kwargs: P.kwargs, # pyright: ignore[reportGeneralTypeIssues] +) -> Array | tuple[Array, ...]: + args_not_none = [arg for arg in args if arg is not None] + array_args = [arg for arg in args_not_none if not is_python_scalar(arg)] + if not array_args: + msg = "Must have at least one argument array" + raise ValueError(msg) + if xp is None: + xp = array_namespace(*args) + + # Normalize and validate shape and dtype + shapes: list[tuple[int | None, ...]] + dtypes: list[DType] + multi_output = False + + if shape is None: + shapes = [np.broadcast_shapes(*(arg.shape for arg in array_args))] * len( + array_args + ) + elif all(isinstance(s, int | None) for s in shape): + # Do not test for shape to be a tuple + # https://github.com/data-apis/array-api/issues/891#issuecomment-2637430522 + shapes = [cast(tuple[int | None, ...], shape)] + else: + shapes = list(shape) # type: ignore[arg-type] # pyright: ignore[reportAssignmentType] + multi_output = True + + if dtype is None: + dtypes = [xp.result_type(*args_not_none)] * len(shapes) + elif multi_output: + if not isinstance(dtype, Sequence): + msg = "Got multiple shapes but only one dtype" + raise ValueError(msg) + dtypes = list(dtype) # pyright: ignore[reportUnknownArgumentType] + else: + if isinstance(dtype, Sequence): + msg = "Got single shape but multiple dtypes" + raise ValueError(msg) + + dtypes = [dtype] + + if len(shapes) != len(dtypes): + msg = f"Got {len(shapes)} shapes and {len(dtypes)} dtypes" + raise ValueError(msg) + del shape + del dtype + # End of shape and dtype parsing + + # Backend-specific branches + if is_dask_namespace(xp): + import dask + import dask.array + + metas: list[Array] = [arg._meta for arg in array_args] # pylint: disable=protected-access # pyright: ignore[reportAttributeAccessIssue] + meta_xp = array_namespace(*metas) + + def dask_wrapper(*args, **kwargs): + while cond(*args, **kwargs): + args = func(*args, **kwargs) + return args + + wrapped = dask.delayed( # type: ignore[attr-defined] # pyright: ignore[reportPrivateImportUsage] + # _lazy_apply_wrapper(dask_wrapper, as_numpy, multi_output, meta_xp), + partial(dask_wrapper, **kwargs), + pure=True, + ) + # This finalizes each arg, which is the same as arg.rechunk(-1). + # Please read docstring above for why we're not using + # dask.array.map_blocks or dask.array.blockwise! + # delayed_out = dask.array.map_blocks( + # wrapped, *args, dtype=dtypes[0], meta=metas[0], **kwargs + # ) + delayed_out = wrapped(*args) + out = tuple( + xp.from_delayed( + delayed_out[i], # pyright: ignore[reportIndexIssue] + # Dask's unknown shapes diverge from the Array API specification + shape=tuple(math.nan if s is None else s for s in shape), + dtype=dtype, + meta=metas[i], + ) + for i, (shape, dtype) in enumerate(zip(shapes, dtypes, strict=True)) + ) + + elif is_jax_namespace(xp) and _is_jax_jit_enabled(xp): + # Delay calling func with jax.pure_callback, which will forward to func eager + # JAX arrays. Do not use jax.pure_callback when running outside of the JIT, + # as it does not support raising exceptions: + # https://github.com/jax-ml/jax/issues/26102 + import jax + + if any(None in shape for shape in shapes): + msg = "Output shape must be fully known when running inside jax.jit" + raise ValueError(msg) + + # Shield kwargs from being coerced into JAX arrays. + # jax.pure_callback calls jax.jit under the hood, but without the chance of + # passing static_argnames / static_argnums. + def jax_func(args, **kwargs): + return func(*args, **kwargs) + + def jax_cond(args, **kwargs): + return cond(*args, **kwargs) + + wrapped = _lazy_apply_wrapper( + partial(jax_func, **kwargs), as_numpy, multi_output, xp + ) + + # suppress unused-ignore to run mypy in -e lint as well as -e dev + out = cast( # type: ignore[bad-cast,unused-ignore] + tuple[Array, ...], + jax.lax.while_loop( + jax_cond, + wrapped, + init_val=args, + ), + ) + + else: + # Eager backends, including non-jitted JAX + # wrapped = _lazy_apply_wrapper(func, as_numpy, multi_output, xp) + while cond(*args): + args = func(*args) + out = args + + return out + + +def traverse_args(args): + if not isinstance(args, Sequence): + yield args + else: + for a in args: + yield from traverse_args(a) + + +def _lazy_apply_wrapper( # type: ignore[explicit-any] # numpydoc ignore=PR01,RT01 + func: Callable[..., Array | ArrayLike | Sequence[Array | ArrayLike]], + as_numpy: bool, + multi_output: bool, + xp: ModuleType, +) -> Callable[..., tuple[Array, ...]]: + """ + Helper of `lazy_apply`. + + Given a function that accepts one or more arrays as positional arguments and returns + a single array-like or a sequence of array-likes, return a function that accepts the + same number of Array API arrays and always returns a tuple of Array API array. + + Any keyword arguments are passed through verbatim to the wrapped function. + """ + + # On Dask, @wraps causes the graph key to contain the wrapped function's name + @wraps(func) + def wrapper( # type: ignore[decorated-any,explicit-any] + *args: Array | complex | None, **kwargs: Any + ) -> tuple[Array, ...]: # numpydoc ignore=GL08 + args_list = [] + device = None + for a in args: + for arg in a: + if arg is not None and not is_python_scalar(arg): + if device is None: + device = _compat.device(arg) + if as_numpy: + import numpy as np + + arg = cast(Array, np.asarray(arg)) # type: ignore[bad-cast] # noqa: PLW2901 + args_list.append(arg) + # assert device is not None + + out = func(tuple(args_list), **kwargs) + + return tuple(xp.asarray(o, device=device) for o in out) + + return wrapper + + +def cross(x1, x2, /, *, axis: int = -1): + xp = array_namespace(x1, x2) + if hasattr(xp, "linalg") and hasattr(xp.linalg, "cross"): + return xp.linalg.cross(x1, x2) + slices = tuple(slice(None, None, None) for i in range(abs(axis + 1))) + c0 = ( + x1[..., 1, *slices] * x2[..., 2, *slices] + - x2[..., 1, *slices] * x1[..., 2, *slices] + ) + c1 = -( + x1[..., 0, *slices] * x2[..., 2, *slices] + - x2[..., 0, *slices] * x1[..., 2, *slices] + ) + c2 = ( + x1[..., 0, *slices] * x2[..., 1, *slices] + - x2[..., 0, *slices] * x1[..., 1, *slices] + ) + return xp.stack((c0, c1, c2), axis=axis) From 04e478c8ff65b9890ada44a6da532dfc6fa47b11 Mon Sep 17 00:00:00 2001 From: purepani Date: Tue, 20 May 2025 16:26:39 -0500 Subject: [PATCH 16/16] Cylinder functions --- src/magpylib/_src/fields/special_cel.py | 39 ++- src/magpylib/_src/fields/special_el3.py | 282 ++++++++++--------- src/magpylib/_src/fields/special_elliptic.py | 1 + 3 files changed, 188 insertions(+), 134 deletions(-) diff --git a/src/magpylib/_src/fields/special_cel.py b/src/magpylib/_src/fields/special_cel.py index 715c5f001..d565d4c2c 100644 --- a/src/magpylib/_src/fields/special_cel.py +++ b/src/magpylib/_src/fields/special_cel.py @@ -7,6 +7,8 @@ import numpy as np from array_api_compat import array_namespace +from magpylib._src.array_api_utils import lazy_while, xp_promote + def cel0(kc, p, c, s): """ @@ -59,6 +61,8 @@ def celv(kc, p, c, s): vectorized version of the cel integral above """ xp = array_namespace(kc, p, c, s) + kc, p, c, s = xp_promote(kc, p, c, s, force_floating=True, xp=xp) + dtype = kc.dtype # if kc == 0: # return NaN @@ -67,7 +71,7 @@ def celv(kc, p, c, s): n = kc.shape[0] k = xp.abs(kc) - em = xp.ones(n, dtype=xp.float64) + em = xp.ones(n, dtype=dtype) # cc = xp.asarray(c, copy=True) # pp = xp.asarray(p, copy=True) @@ -106,10 +110,8 @@ def celv(kc, p, c, s): em = k + em kk = xp.asarray(k, copy=True) - # define a mask that adjusts with every evaluation step so that only - # non-converged entries are further iterated. - mask = xp.ones(n, dtype=xp.bool) - while xp.any(mask): + def iter_step(f, cc, g, ss, pp, em, kk, k): + # f, cc, g, ss, pp, em, kk, k = val k = 2 * xp.sqrt(kk) kk = k * em f = cc @@ -119,11 +121,22 @@ def celv(kc, p, c, s): pp = g + pp g = em em = k + em + return f, cc, g, ss, pp, em, kk, k - # redefine mask + def cond(f, cc, g, ss, pp, em, kk, k): + # f, cc, g, ss, pp, em, kk, k = val err = g - k tol = g * errtol mask = xp.abs(err) > tol + return xp.any(mask) + + # define a mask that adjusts with every evaluation step so that only + # non-converged entries are further iterated. + mask = xp.ones(n, dtype=xp.bool) + val = f, cc, g, ss, pp, em, kk, k + val = lazy_while(iter_step, cond, val) + + f, cc, g, ss, pp, em, kk, k = val return (xp.pi / 2) * (ss + cc * em) / (em * (em + pp)) @@ -191,7 +204,9 @@ def cel_iterv(qc, p, g, cc, ss, em, kk): Iterative part of Bulirsch cel algorithm """ xp = array_namespace(qc, p, g, cc, ss, em, kk) - while xp.any(xp.abs(g - qc) >= qc * 1e-8): + + def iter_step(qc, p, g, cc, ss, em, kk): + # qc, p, g, cc, ss, em, kk = val qc = 2 * xp.sqrt(kk) kk = qc * em f = cc @@ -201,4 +216,14 @@ def cel_iterv(qc, p, g, cc, ss, em, kk): p = p + g g = em em = em + qc + return qc, p, g, cc, ss, em, kk + + def cond(qc, p, g, cc, ss, em, kk): + # qc, p, g, cc, ss, em, kk = val + return xp.any(xp.abs(g - qc) >= qc * 1e-8) + + val = qc, p, g, cc, ss, em, kk + val = lazy_while(iter_step, cond, val) + print(val) + qc, p, g, cc, ss, em, kk = val return 1.5707963267948966 * (ss + cc * em) / (em * (em + p)) diff --git a/src/magpylib/_src/fields/special_el3.py b/src/magpylib/_src/fields/special_el3.py index 8d1d35298..0cc9e6361 100644 --- a/src/magpylib/_src/fields/special_el3.py +++ b/src/magpylib/_src/fields/special_el3.py @@ -1,5 +1,6 @@ from __future__ import annotations +import array_api_extra as xpx import numpy as np from array_api_compat import array_namespace @@ -273,8 +274,12 @@ def el3v(x, kc, p): s = xp.zeros(nnn) mask1 = kc == 0 - s[mask1] = CA / (1 + xp.abs(x[mask1])) - s[~mask1] = kc[~mask1] + s = xpx.apply_where( + mask1, + (x, kc), + lambda x, kc: CA / (1 + xp.abs(x)), + lambda x, kc: kc, + ) t = s * s pm = 0.5 * t e = hh * t @@ -316,7 +321,9 @@ def el3v(x, kc, p): if bo: sx = -sx ux = (ux + 1) * 0.5 - result[mask2] = (ux - sx * hx) * xp.sqrt(hx) * xx + ux * xp.asinh(xx) + result = xpx.at(result)[mask2].set( + (ux - sx * hx) * xp.sqrt(hx) * xx + ux * xp.asinh(xx) + ) mask2x = ~mask2 p, pm, t, hh, h, x = ( @@ -339,14 +346,14 @@ def el3v(x, kc, p): p1 = xp.asarray(p, copy=True) mask3 = p == 0 - p1[mask3] = CB / hh[mask3] + p1 = xpx.apply_where(mask3, (hh, p1), lambda hh, p1: CB / hh, lambda hh, p1: p1) s = xp.abs(s) y = xp.abs(x) g = p1 - 1.0 - g[g == 0] = CB + g = xpx.at(g)[g == 0].set(CB) f = p1 - t mask4 = f == 0 - f[mask4] = CB * t[mask4] + f = xp.where(mask4, CB * t, f) am = 1.0 - t ap = 1.0 + e r = p1 * h @@ -357,7 +364,7 @@ def el3v(x, kc, p): de = xp.sqrt(pz) q = xp.sqrt(xp.abs(p1)) mask5 = pm > 0.5 - pm[mask5] = 0.5 + pm = xpx.at(pm)[mask5].set(0.5) pm = p1 - pm uvdc = xp.zeros((4, (pm.shape[0]))) @@ -365,77 +372,79 @@ def el3v(x, kc, p): mask6 = pm >= 0.0 if xp.any(mask6): - u[mask6] = xp.sqrt(r[mask6] * ap[mask6]) - v[mask6] = y[mask6] * de[mask6] * xp.sign(g[mask6]) - d[mask6] = 1 / q[mask6] - c[mask6] = 1.0 + u = xp.where(mask6, xp.sqrt(r * ap), u) + v = xp.where(mask6, y * de * xp.sign(g), v) + d = xp.where(mask6, 1 / q, d) + c = xp.where(mask6, 1.0, c) mask6x = ~mask6 if xp.any(mask6x): - u[mask6x] = xp.sqrt(h[mask6x] * ap[mask6x] * pz[mask6x]) - ye[mask6x] = y[mask6x] * q[mask6x] - v[mask6x] = am[mask6x] * ye[mask6x] - q[mask6x] = -de[mask6x] / g[mask6x] - d[mask6x] = -am[mask6x] / de[mask6x] - c[mask6x] = 0 - pz[mask6x] = ap[mask6x] - r[mask6x] + u = xp.where(mask6x, xp.sqrt(h * ap * pz), u) + ye = xp.where(mask6x, y * q, ye) + v = xp.where(mask6x, am * ye, v) + q = xp.where(mask6x, -de / g, q) + d = xp.where(mask6x, -am / de, d) + c = xpx.at(c)[mask6x].set(0) + pz = xp.where(mask6x, ap - r, pz) if xp.any(bo): - r[bo] = v[bo] / u[bo] - z[bo] = 1 - k[bo] = 1 + r = xp.where(bo, v / u, r) + z = xpx.at(z)[bo].set(1) + k = xpx.at(k)[bo].set(1) mask7 = bo & (pm < 0) if xp.any(mask7): - h[mask7] = y[mask7] * xp.sqrt(h[mask7] / (ap[mask7] * fa[mask7])) - h[mask7] = 1 / h[mask7] - h[mask7] - z[mask7] = h[mask7] - 2 * r[mask7] - r[mask7] = 2 + r[mask7] * h[mask7] + h = xp.where(mask7, y * xp.sqrt(h / (ap * fa)), h) + h = xp.where(mask7, 1 / h - h, h) + z = xp.where(mask7, h - 2 * r, z) + r = xp.where(mask7, 2 + r * h, r) mask7a = mask7 & (r == 0) - r[mask7a] = CB + r = xpx.at(r)[mask7a].set(CB) mask7b = mask7 & (z == 0) - z[mask7b] = h[mask7b] * CB + z = xp.where(mask7b, h * CB, z) - z[mask7] = r[mask7] / z[mask7] - r[mask7] = xp.asarray(z[mask7], copy=True) - w[mask7] = pz[mask7] + z = xp.where(mask7, r / z, z) + r = xp.where(mask7, z, r) + w = xp.where(mask7, pz, w) - u[bo] = u[bo] / w[bo] - v[bo] = v[bo] / w[bo] + u = xp.where(bo, u / w, u) + v = xp.where(bo, v / w, v) box = ~bo if xp.any(box): - t[box] = u[box] + xp.abs(v[box]) - bk[box] = True + t = xp.where(box, u + xp.abs(v), t) + bk = xpx.at(bk)[box].set(True) mask8 = box & (p1 < 0) if xp.any(mask8): - de[mask8] = v[mask8] / pz[mask8] - ye[mask8] = u[mask8] * ye[mask8] - ye[mask8] = 2 * ye[mask8] - u[mask8] = t[mask8] / pz[mask8] - v[mask8] = (-f[mask8] - g[mask8] * e[mask8]) / t[mask8] - t[mask8] = pz[mask8] * xp.abs(w[mask8]) - z[mask8] = ( - hh[mask8] * r[mask8] * f[mask8] - g[mask8] * ap[mask8] + ye[mask8] - ) / t[mask8] - ye[mask8] = ye[mask8] / t[mask8] + de = xp.where(mask8, v / pz, de) + ye = xp.where(mask8, u * ye, ye) + ye = xp.where(mask8, 2 * ye, ye) + u = xp.where(mask8, t / pz, u) + v = xp.where(mask8, (-f - g * e) / t, v) + t = xp.where(mask8, pz * xp.abs(w), t) + z = xp.where( + mask8, + (hh * r * f - g * ap + ye) / t, + z, + ) + ye = xp.where(mask8, ye / t, ye) mask8x = box & (p1 >= 0) if xp.any(mask8x): - de[mask8x] = v[mask8x] / w[mask8x] - ye[mask8x] = 0 - u[mask8x] = (e[mask8x] + p1[mask8x]) / t[mask8x] - v[mask8x] = t[mask8x] / w[mask8x] - z[mask8x] = 1.0 + de = xp.where(mask8x, v / w, de) + ye = xpx.at(ye).set(0) + u = xp.where(mask8x, (e + p1) / t, u) + v = xp.where(mask8x, t / w, v) + z = xpx.at(z)[mask8x].set(1.0) mask9 = box & (s > 1) if xp.any(mask9): - h[mask9] = u[mask9] - u[mask9] = v[mask9] - v[mask9] = h[mask9] + h = xp.where(mask9, u, h) + u = xp.where(mask9, v, u) + v = xp.where(mask9, h, v) y = 1 / y e = xp.asarray(s, copy=True) @@ -448,120 +457,118 @@ def el3v(x, kc, p): (p.shape[0]), dtype=xp.bool ) # dynamic mask, changed in each loop iteration while xp.any(mask10): - y[mask10] = y[mask10] - e[mask10] / y[mask10] + y = xp.where(mask10, y - e / y, y) mask11 = mask10 & (y == 0.0) - y[mask11] = xp.sqrt(e[mask11]) * CB - - f[mask10] = c[mask10] - c[mask10] = d[mask10] / q[mask10] + c[mask10] - g[mask10] = e[mask10] / q[mask10] - d[mask10] = f[mask10] * g[mask10] + d[mask10] - d[mask10] = 2 * d[mask10] - q[mask10] = g[mask10] + q[mask10] - g[mask10] = t[mask10] - t[mask10] = s[mask10] + t[mask10] - n[mask10] = 2 * n[mask10] - m[mask10] = 2 * m[mask10] + y = xp.where(mask11, xp.sqrt(e) * CB, y) + + f = xp.where(mask10, c, f) + c = xp.where(mask10, d / q + c, c) + g = xp.where(mask10, e / q, g) + d = xp.where(mask10, f * g + d, d) + d = xp.where(mask10, 2 * d, d) + q = xp.where(mask10, g + q, q) + g = xp.where(mask10, t, g) + t = xp.where(mask10, s + t, t) + n = xp.where(mask10, 2 * n, n) + m = xp.where(mask10, 2 * m, m) bo10 = mask10 & bo if xp.any(bo10): bo10b = bo10 & (z < 0) - m[bo10b] = k[bo10b] + m[bo10b] + m = xp.where(bo10b, k + m, m) - k[bo10] = xp.sign(r[bo10]) - h[bo10] = e[bo10] / (u[bo10] * u[bo10] + v[bo10] * v[bo10]) - u[bo10] = u[bo10] * (1 + h[bo10]) - v[bo10] = v[bo10] * (1 - h[bo10]) + k = xp.where(bo10, xp.sign(r), k) + h = xp.where(bo10, e / (u * u + v * v), h) + u = xp.where(bo10, u * (1 + h), u) + v = xp.where(bo10, v * (1 - h), v) bo10x = xp.asarray(mask10 & ~bo10, dtype=xp.bool) if xp.any(bo10x): - r[bo10x] = u[bo10x] / v[bo10x] - h[bo10x] = z[bo10x] * r[bo10x] - z[bo10x] = h[bo10x] * z[bo10x] - hh[bo10x] = e[bo10x] / v[bo10x] + r = xp.where(bo10x, u / v, r) + h = xp.where(bo10x, z * r, h) + z = xp.where(bo10x, h * z, z) + hh = xp.where(bo10x, e / v, hh) bo10x_bk = xp.asarray(bo10x * bk, dtype=bool) # if bk bo10x_bkx = xp.asarray(bo10x * ~bk, dtype=bool) if xp.any(bo10x_bk): - de[bo10x_bk] = de[bo10x_bk] / u[bo10x_bk] - ye[bo10x_bk] = ye[bo10x_bk] * (h[bo10x_bk] + 1 / h[bo10x_bk]) + de[ - bo10x_bk - ] * (1 + r[bo10x_bk]) - de[bo10x_bk] = de[bo10x_bk] * (u[bo10x_bk] - hh[bo10x_bk]) - bk[bo10x_bk] = xp.abs(ye[bo10x_bk]) < 1 + de = xp.where(bo10x_bk, de / u, de) + ye = xp.where(bo10x_bk, ye * (h + 1 / h) + de * (1 + r), ye) + de = xp.where(bo10x_bk, de[bo10x_bk] * (u[bo10x_bk] - hh[bo10x_bk]), de) + bk = xp.where(bo10x_bk, xp.abs(ye[bo10x_bk]) < 1, bk) if xp.any(bo10x_bkx): - a_crack = xp.log(x[bo10x_bkx]) - k[bo10x_bkx] = (a_crack / ln2).astype(int) + 1 - a_crack = a_crack - k[bo10x_bkx] * ln2 - m[bo10x_bkx] = xp.exp(a_crack) - m[bo10x_bkx] = m[bo10x_bkx] + k[bo10x_bkx] + a_crack = xp.log(x) + k = xp.where(bo10x_bkx, xp.astype(a_crack / ln2, xp.int32) + 1, k) + a_crack = a_crack - k * ln2 + m = xp.where(bo10x_bkx, xp.exp(a_crack), m) + m = xp.where(bo10x_bkx, m + k, m) mask11 = xp.abs(g - s) > CA * g if xp.any(mask11): bo11 = mask11 & bo if xp.any(bo11): - g[bo11] = (1 / r[bo11] - r[bo11]) * 0.5 - hh[bo11] = u[bo11] + v[bo11] * g[bo11] - h[bo11] = g[bo11] * u[bo11] - v[bo11] + g = xp.where(bo11, (1 / r - r) * 0.5, g) + hh = xp.where(bo11, u + v * g, hh) + h = xp.where(bo11, g * u - v, h) bo11b = bo11 & (hh == 0) - hh[bo11b] = u[bo11b] * CB + hh = xp.where(bo11b, u * CB, hh) bo11c = bo11 & (h == 0) - h[bo11c] = v[bo11c] * CB + h = xp.where(bo11c, v * CB, h) - z[bo11] = r[bo11] * h[bo11] - r[bo11] = hh[bo11] / h[bo11] + z = xp.where(bo11, r * h, z) + r = xp.where(bo11, hh / h, r) bo11x = mask11 & ~bo if xp.any(bo11x): - u[bo11x] = u[bo11x] + e[bo11x] / u[bo11x] - v[bo11x] = v[bo11x] + hh[bo11x] + u = xp.where(bo11x, u + e / u, u) + v = xp.where(bo11x, v + hh, v) - s[mask11] = xp.sqrt(e[mask11]) - s[mask11] = 2 * s[mask11] - e[mask11] = s[mask11] * t[mask11] - l[mask11] = 2 * l[mask11] + s = xp.where(mask11, xp.sqrt(e), s) + s = xp.where(mask11, 2 * s, s) + e = xp.where(mask11, s * t, e) + l = xp.where(mask11, 2 * l, l) mask12 = mask11 & (y < 0) - l[mask12] = l[mask12] + 1 + l = xp.where(mask12, l + 1, l) # break off parts that have completed their iteration mask10 = mask11 mask12 = y < 0 - l[mask12] = l[mask12] + 1 + l = xp.where(mask12, l + 1, l) e = xp.atan(t / y) + xp.pi * l e = e * (c * t + d) / (t * (t + q)) if xp.any(bo): - h[bo] = v[bo] / (t[bo] + u[bo]) - z[bo] = 1 - r[bo] * h[bo] - h[bo] = r[bo] + h[bo] + h = xp.where(bo, v / (t + u), h) + z = xp.where(bo, 1 - r * h, z) + h = xp.where(bo, r + h, h) bob = bo & (z == 0) - z[bob] = CB + z = xp.where(bob, CB, z) boc = bo & (z < 0) - m[boc] = m[boc] + xp.sign(h[boc]) + m = xp.where(boc, m + xp.sign(h), m) - s[bo] = xp.atan(h[bo] / z[bo]) + m[bo] * xp.pi + s = xp.where(bo, xp.atan(h / z) + m * xp.pi, s) box = ~bo if xp.any(box): box_bk = box & bk - s[box_bk] = xp.asinh(ye[box_bk]) + s = xp.where(box_bk, xp.asinh(ye), s) box_bkx = box & ~bk - s[box_bkx] = xp.log(z[box_bkx]) + m[box_bkx] * ln2 + s = xp.where(box_bkx, xp.log(z) + m * ln2, s) - s[box] = s[box] * 0.5 + s = xp.where(box, s * 0.5, s) e = (e + xp.sqrt(fa) * s) / n - result[~mask2] = xp.sign(x) * e + result = xpx.at(result)[~mask2].set(xp.sign(x) * e) # include mask0-case - result0[mask0] = result + result0 = xpx.at(result0)[mask0].set(result) return result0 @@ -608,12 +615,12 @@ def el3_angle(phi: np.ndarray, n: np.ndarray, m: np.ndarray) -> np.ndarray: mask1 = (n <= 0) & (phi_red < -xp.pi / 2) mask2 = (n >= 0) & (phi_red > xp.pi / 2) if xp.any(mask1): - n[mask1] = n[mask1] - 1 - phi_red[mask1] = phi_red[mask1] + xp.pi + n = xp.where(mask1, n - 1, n) + phi_red = xp.where(mask1, phi_red + xp.pi, phi_red) if xp.any(mask2): - n[mask2] = n[mask2] + 1 - phi_red[mask2] = phi_red[mask2] - xp.pi + n = xp.where(mask2, n + 1, n) + phi_red = xp.where(mask2, phi_red - xp.pi, phi_red) mask3 = n != 0 mask3x = ~mask3 @@ -630,14 +637,26 @@ def el3_angle(phi: np.ndarray, n: np.ndarray, m: np.ndarray) -> np.ndarray: mask3c = ~mask3a & ~mask3b if xp.any(mask3a): - results3[mask3a] = (2 * n3[mask3a] + 1) * cel3_res[mask3a] + results3 = xp.where(mask3a, (2 * n3 + 1) * cel3_res, results3) if xp.any(mask3b): - results3[mask3b] = (2 * n3[mask3b] - 1) * cel3_res[mask3b] + results3 = xp.where(mask3b, (2 * n3 - 1) * cel3_res, results3) if xp.any(mask3c): - el33_res = el3(xp.tan(phi3[mask3c]), kc3[mask3c], p3[mask3c]) - results3[mask3c] = 2 * n3[mask3c] * cel3_res[mask3c] + el33_res + results3 = xpx.apply_where( + mask3c, + ( + phi3, + kc3, + p3, + n3, + cel3_res, + results3, + ), + lambda phi3, kc3, p3, n3, cel3_res, results3: el3(xp.tan(phi3), kc3, p3) + + 2 * n3 * cel3_res, + lambda phi3, kc3, p3, n3, cel3_res, results3: results3, + ) - results[mask3] = results3 + results = xpx.at(results)[mask3].set(results3) if xp.any(mask3x): phi_red3x = phi_red[mask3x] @@ -650,19 +669,28 @@ def el3_angle(phi: np.ndarray, n: np.ndarray, m: np.ndarray) -> np.ndarray: if xp.any(mask3xa): onez = xp.ones(xp.count_nonzero(mask3xa)) - results3x[mask3xa] = cel( - kc3x[mask3xa], p3x[mask3xa], onez, onez + results3x = xpx.apply_where( + mask3xa, + (kc3x, p3x, results3x), + lambda kc3x, p3x, results3x: cel(kc3x, p3x, onez, onez), + lambda kc3x, p3x, results3x: results3x, ) # 3rd kind cel if xp.any(mask3xb): onez = xp.ones(xp.count_nonzero(mask3xb)) - results3x[mask3xb] = -cel( - kc3x[mask3xb], p3x[mask3xb], onez, onez + results3x = xp.where( + mask3xb, + (kc3x, p3x, results3x), + lambda kc3x, p3x, results3x: -cel(kc3x, p3x, onez, onez), + lambda kc3x, p3x, results3x: results3x, ) # 3rd kind cel if xp.any(mask3xc): - results3x[mask3xc] = el3( - xp.tan(phi3x[mask3xc]), kc3x[mask3xc], p3x[mask3xc] + results3x = xpx.apply_where( + mask3xc, + (phi3x, kc3x, p3x, results3x), + lambda phi3x, kc3x, p3x, results3x: el3(xp.tan(phi3x), kc3x, p3x), + lambda phi3x, kc3x, p3x, results3x: results3x, ) - results[mask3x] = results3x + results = xpx.at(results)[mask3x].set(results3x) return results diff --git a/src/magpylib/_src/fields/special_elliptic.py b/src/magpylib/_src/fields/special_elliptic.py index 6aad39794..c2ef9ccf7 100644 --- a/src/magpylib/_src/fields/special_elliptic.py +++ b/src/magpylib/_src/fields/special_elliptic.py @@ -17,6 +17,7 @@ def ellipkinc(phi: Array, m: Array): def ellipe(m: Array): + # return cel(kc, one, one, kc2) return xpx.lazy_apply(sp.special.ellipe, m, as_numpy=True)