added unittests for coordinates module

balltree/__init__.py

@@ -1,7 +1,7 @@
-__version__ = "0.1"
+__version__ = "0.2"
 
-from ._balltree import BallTree, default_leafsize
 from .angulartree import AngularTree
+from .balltree import BallTree, default_leafsize
 
 __all__ = [
     "AngularTree",

balltree/angulartree.py

@@ -4,7 +4,11 @@
 from numpy.typing import NDArray
 
 from balltree import coordinates as coord
-from balltree._balltree import BallTree, default_leafsize
+from balltree.balltree import BallTree, default_leafsize
+
+__all__ = [
+    "AngularTree",
+]
 
 
 class AngularTree:

balltree/balltree.c

@@ -1,3 +1,4 @@
+#define PY_SSIZE_T_CLEAN
 #include <Python.h>
 #include <numpy/arrayobject.h>
 
@@ -934,12 +935,12 @@ static PyTypeObject PyBallTreeType = {
 
 static struct PyModuleDef pyballtree = {
     PyModuleDef_HEAD_INIT,
-    .m_name = "balltree._balltree",
+    .m_name = "balltree.balltree",
     .m_doc = "Fast balltree implementation",
     .m_size = -1,
 };
 
-PyMODINIT_FUNC PyInit__balltree(void) {
+PyMODINIT_FUNC PyInit_balltree(void) {
     if (PyType_Ready(&PyBallTreeType) < 0) {
         return NULL;
     }

balltree/coordinates.py

@@ -3,9 +3,19 @@
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
+__all__ = [
+    "sgn",
+    "angular_to_cylinder",
+    "cylinder_to_angular",
+    "angular_to_euclidean",
+    "euclidean_to_angular",
+    "angle_to_chorddist",
+    "chorddist_to_angle",
+]
+
 
 def sgn(val: ArrayLike) -> NDArray[np.float64]:
-    return np.where(val == 0, 1.0, np.sign(val, dtype=np.float64))
+    return 2.0 * (val >= 0.0) - 1.0  # 1.0 if (val >= 0.0) else -1.0
 
 
 def angular_to_cylinder(radec: NDArray) -> NDArray[np.float64]:
@@ -48,7 +58,7 @@ def euclidean_to_angular(xyz: NDArray) -> NDArray[np.float64]:
     x_normed = np.ones_like(x)  # fallback for zero-division, arccos(1)=0.0
     np.divide(x, r_d2, where=r_d2 > 0.0, out=x_normed)
     radec[:, 0] = np.arccos(x_normed) * sgn(y) % (2.0 * np.pi)
-    radec[:, 1] = np.arcsin(x / r_d3)
+    radec[:, 1] = np.arcsin(z / r_d3)
     return radec
 
 

setup.py

@@ -2,7 +2,7 @@
 from setuptools import Extension, setup
 
 balltree = Extension(
-    name="balltree._balltree",
+    name="balltree.balltree",
     sources=[
         "balltree/balltree.c",
         "src/point.c",

tests/test_coordinates.py

@@ -5,37 +5,150 @@
 from balltree import coordinates
 
 
+@pytest.fixture
+def coords_angular_ra():
+    return np.deg2rad(
+        [
+            [-90.0, 0.0],
+            [0.0, 0.0],
+            [90.0, 0.0],
+            [180.0, 0.0],
+            [270.0, 0.0],
+            [360.0, 0.0],
+            [450.0, 0.0],
+        ]
+    )
+
+
+@pytest.fixture
+def coords_cylinder_ra(coords_angular_ra):
+    return coords_angular_ra
+
+
+@pytest.fixture
+def coords_euclidean_ra():
+    return np.array(
+        [
+            [0.0, -1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [-1.0, 0.0, 0.0],
+            [0.0, -1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+        ]
+    )
+
+
+@pytest.fixture
+def coords_angular_dec():
+    return np.deg2rad(
+        [
+            [-90.0, 90.0],
+            [0.0, 90.0],
+            [90.0, 90.0],
+            [-90.0, -90.0],
+            [0.0, -90.0],
+            [90.0, -90.0],
+        ]
+    )
+
+
+@pytest.fixture
+def coords_cylinder_dec(coords_angular_dec):
+    values = coords_angular_dec.copy()
+    values[:, 1] = coordinates.sgn(coords_angular_dec[:, 1])
+    return values
+
+
+@pytest.fixture
+def coords_euclidean_dec():
+    return np.array(
+        [
+            [0.0, 0.0, 1.0],
+            [0.0, 0.0, 1.0],
+            [0.0, 0.0, 1.0],
+            [0.0, 0.0, -1.0],
+            [0.0, 0.0, -1.0],
+            [0.0, 0.0, -1.0],
+        ]
+    )
+
+
+@pytest.fixture
+def chord_length():
+    return np.array([0.0, np.sqrt(2.0), 2.0])
+
+
+@pytest.fixture
+def angle():
+    return np.array([0.0, np.pi / 2.0, np.pi])
+
+
 def test_sgn():
     values = np.array([-1.0, -0.5, 0.0, 0.5, 1.0])
     result = np.array([-1.0, -1.0, 1.0, 1.0, 1.0])
     npt.assert_array_equal(coordinates.sgn(values), result)
 
 
-@pytest.mark.skip
-def test_angular_to_cylinder():
-    pass
+def test_angular_to_cylinder_ra(coords_angular_ra, coords_cylinder_ra):
+    result = coordinates.angular_to_cylinder(coords_angular_ra)
+    expected = coords_cylinder_ra
+    npt.assert_array_almost_equal(result, expected)
+
+
+def test_angular_to_cylinder_dec(coords_angular_dec, coords_cylinder_dec):
+    result = coordinates.angular_to_cylinder(coords_angular_dec)
+    expected = coords_cylinder_dec
+    npt.assert_array_almost_equal(result, expected)
+
+
+def test_cylinder_to_angular_ra(coords_cylinder_ra, coords_angular_ra):
+    result = coordinates.cylinder_to_angular(coords_cylinder_ra)
+    expected = coords_angular_ra
+    npt.assert_array_almost_equal(result, expected)
+
+
+def test_cylinder_to_angular_dec(coords_cylinder_dec, coords_angular_dec):
+    result = coordinates.cylinder_to_angular(coords_cylinder_dec)
+    expected = coords_angular_dec
+    npt.assert_array_almost_equal(result, expected)
+
+
+def test_angular_to_euclidean_ra(coords_angular_ra, coords_euclidean_ra):
+    result = coordinates.angular_to_euclidean(coords_angular_ra)
+    expected = coords_euclidean_ra
+    npt.assert_array_almost_equal(result, expected)
 
 
-@pytest.mark.skip
-def test_cylinder_to_angular():
-    pass
+def test_angular_to_euclidean_dec(coords_angular_dec, coords_euclidean_dec):
+    result = coordinates.angular_to_euclidean(coords_angular_dec)
+    expected = coords_euclidean_dec
+    npt.assert_array_almost_equal(result, expected)
 
 
-@pytest.mark.skip
-def test_angular_to_euclidean():
-    pass
+def test_euclidean_to_angular_ra(coords_euclidean_ra, coords_angular_ra):
+    result = coordinates.euclidean_to_angular(coords_euclidean_ra)
+    # some points are outside of the range of [0, 2pi)
+    expected = coords_angular_ra % (2.0 * np.pi)
+    npt.assert_array_almost_equal(result, expected)
 
 
-@pytest.mark.skip
-def test_euclidean_to_angular():
-    pass
+def test_euclidean_to_angular_dec(coords_euclidean_dec, coords_angular_dec):
+    result = coordinates.euclidean_to_angular(coords_euclidean_dec)
+    # all tested points are coordinate singularities, disregard RA component
+    expected = coords_angular_dec.copy()
+    expected[:, 0] = 0.0
+    npt.assert_array_almost_equal(result, expected)
 
 
-@pytest.mark.skip
-def test_angle_to_chorddist():
-    pass
+def test_angle_to_chorddist(angle, chord_length):
+    result = coordinates.angle_to_chorddist(angle)
+    expected = chord_length
+    npt.assert_array_almost_equal(result, expected)
 
 
-@pytest.mark.skip
-def test_chorddist_to_angle():
-    pass
+def test_chorddist_to_angle(chord_length, angle):
+    result = coordinates.chorddist_to_angle(chord_length)
+    expected = angle
+    npt.assert_array_almost_equal(result, expected)