fix: fix special._bessel not working (#60)

Author: 34j

src/ultrasphere/special/_bessel.py

@@ -61,17 +61,21 @@ def szv(
             zv = yv
     elif type == "h1":
         if derivative:
-            zv = hankel1
-        else:
             raise AssertionError()
+        else:
+            zv = hankel1
     elif type == "h2":
         if derivative:
-            zv = hankel2
-        else:
             raise AssertionError()
+        else:
+            zv = hankel2
+
+    dtype = xp.result_type(v, d, z)
+    if type in ("h1", "h2"):
+        dtype = xp.result_type(dtype, xp.complex64)
     return (
-        xp.sqrt(xp.pi / 2)
-        * xp.asarray(zv(v + d_half_minus_1, z), device=z.device, dtype=z.dtype)
+        xp.sqrt(xp.asarray(xp.pi / 2, device=z.device, dtype=dtype))
+        * xp.asarray(zv(v + d_half_minus_1, z), device=z.device, dtype=dtype)
         / (z**d_half_minus_1)
     )
 

tests/test_special.py

@@ -1,35 +1,61 @@
 from typing import Literal
 
+import array_api_extra as xpx
 import numpy as np
 import pytest
-from numpy.testing import assert_allclose
-from scipy.special import jv, jvp, spherical_jn, spherical_yn, yv, yvp
+from array_api._2024_12 import ArrayNamespaceFull
+from scipy.special import hankel1, hankel2, jv, jvp, spherical_jn, spherical_yn, yv, yvp
 
-from ultrasphere.special._bessel import sjv, syv
+from ultrasphere.special._bessel import shn1, shn2, sjv, syv
 
 
 @pytest.mark.parametrize("derivative", [True, False])
 @pytest.mark.parametrize("d", [2, 3])
-@pytest.mark.parametrize("type", ["j", "y"])
-def test_sjyn(d: Literal[2, 3], type: Literal["j", "y"], derivative: bool) -> None:
+@pytest.mark.parametrize("type", ["j", "y", "h1", "h2"])
+def test_sjyn(
+    d: Literal[2, 3],
+    type: Literal["j", "y", "h1", "h2"],
+    derivative: bool,
+    xp: ArrayNamespaceFull,
+) -> None:
     n = np.random.randint(5, size=(10, 10))
     z = np.random.random((10, 10))
     if type == "j":
-        actual = sjv(n, np.array(d), z, derivative=derivative)
         if d == 2:
             expected = np.sqrt(np.pi / 2) * (jvp(n, z) if derivative else jv(n, z))
         elif d == 3:
             expected = spherical_jn(n, z, derivative=derivative)
         else:
             raise ValueError("d must be 2 or 3")
+        expected = xp.asarray(expected)
+        actual = sjv(xp.asarray(n), xp.asarray(d), xp.asarray(z), derivative=derivative)
     elif type == "y":
-        actual = syv(n, np.array(d), z, derivative=derivative)
         if d == 2:
             expected = np.sqrt(np.pi / 2) * (yvp(n, z) if derivative else yv(n, z))
         elif d == 3:
             expected = spherical_yn(n, z, derivative=derivative)
         else:
             raise ValueError("d must be 2 or 3")
+        expected = xp.asarray(expected)
+        actual = syv(xp.asarray(n), xp.asarray(d), xp.asarray(z), derivative=derivative)
+    elif type == "h1":
+        if derivative:
+            pytest.skip("derivative of hankel1 not implemented")
+        if d == 2:
+            expected = np.sqrt(np.pi / 2) * hankel1(n, z)
+        elif d == 3:
+            expected = spherical_jn(n, z) + 1j * spherical_yn(n, z)
+        expected = xp.asarray(expected)
+        actual = shn1(xp.asarray(n), xp.asarray(d), xp.asarray(z))
+    elif type == "h2":
+        if derivative:
+            pytest.skip("derivative of hankel2 not implemented")
+        if d == 2:
+            expected = np.sqrt(np.pi / 2) * hankel2(n, z)
+        elif d == 3:
+            expected = spherical_jn(n, z) - 1j * spherical_yn(n, z)
+        expected = xp.asarray(expected)
+        actual = shn2(xp.asarray(n), xp.asarray(d), xp.asarray(z))
     else:
         raise ValueError("type must be 'j' or 'y'")
-    assert_allclose(actual, expected)
+    assert xp.all(xpx.isclose(actual, expected, rtol=1e-6, atol=1e-6))