diff --git a/balltree/__init__.py b/balltree/__init__.py index 97c05da..1e82dfd 100644 --- a/balltree/__init__.py +++ b/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", diff --git a/balltree/angulartree.py b/balltree/angulartree.py index 48bd580..2969d04 100644 --- a/balltree/angulartree.py +++ b/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: diff --git a/balltree/balltree.c b/balltree/balltree.c index 73f982e..a89911a 100644 --- a/balltree/balltree.c +++ b/balltree/balltree.c @@ -1,3 +1,4 @@ +#define PY_SSIZE_T_CLEAN #include #include @@ -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; } diff --git a/balltree/coordinates.py b/balltree/coordinates.py index 178258b..de6d165 100644 --- a/balltree/coordinates.py +++ b/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 diff --git a/setup.py b/setup.py index 75f6847..3a6c2b0 100644 --- a/setup.py +++ b/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", diff --git a/tests/test_coordinates.py b/tests/test_coordinates.py index f20c66f..ae08f57 100644 --- a/tests/test_coordinates.py +++ b/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)