drafted missing tests, created coordinates module

balltree/angulartree.py

@@ -1,83 +1,28 @@
 from __future__ import annotations
 
 import numpy as np
-from numpy.typing import ArrayLike, NDArray
+from numpy.typing import NDArray
 
-from ._balltree import BallTree, default_leafsize
-
-
-def sgn(val: ArrayLike) -> ArrayLike:
-    return np.where(val == 0, 1.0, np.sign(val))
-
-
-def radec_to_xy(radec: NDArray[np.float64]) -> NDArray[np.float64]:
-    radec = np.atleast_2d(radec)
-    xy = np.empty_like(radec)
-    xy[:, 0] = radec[:, 0]
-    xy[:, 1] = np.sin(radec[:, 1])
-    return xy
-
-
-def xy_to_radec(xy: NDArray[np.float64]) -> NDArray[np.float64]:
-    radec = np.empty_like(xy)
-    radec[:, 0] = xy[:, 0]
-    radec[:, 1] = np.arcsin(xy[:, 1])
-    return radec
-
-
-def radec_to_xyz(radec: NDArray[np.float64]) -> NDArray[np.float64]:
-    radec = np.atleast_2d(radec)
-    ra = radec[:, 0]
-    dec = radec[:, 1]
-    cos_dec = np.cos(dec)
-
-    xyz = np.empty((len(ra), 3))
-    xyz[:, 0] = np.cos(ra) * cos_dec
-    xyz[:, 1] = np.sin(ra) * cos_dec
-    xyz[:, 2] = np.sin(dec)
-    return xyz
-
-
-def xyz_to_radec(xyz: NDArray[np.float64]) -> NDArray[np.float64]:
-    xyz = np.atleast_2d(xyz)
-    x = xyz[:, 0]
-    y = xyz[:, 1]
-    z = xyz[:, 2]
-    r_d2 = np.sqrt(x * x + y * y)
-    r_d3 = np.sqrt(x * x + y * y + z * z)
-
-    radec = np.empty((len(x), 2))
-    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)
-    return radec
-
-
-def angle_to_radius(angle: float) -> float:
-    return 2.0 * np.sin(angle / 2.0)
-
-
-def radius_to_angle(radius: float) -> float:
-    return 2.0 * np.arcsin(radius / 2.0)
+from balltree import coordinates as coord
+from balltree._balltree import BallTree, default_leafsize
 
 
 class AngularTree:
     _tree: BallTree
 
     def __init__(
         self,
-        radec: NDArray[np.float64],
-        weight: NDArray[np.float64] | None = None,
+        radec: NDArray,
+        weight: NDArray | None = None,
         leafsize: int = default_leafsize,
     ) -> None:
-        xyz = radec_to_xyz(radec)
+        xyz = coord.radec_to_xyz(radec)
         self._tree = BallTree(xyz, weight, leafsize=leafsize)
 
     @property
     def data(self) -> NDArray[np.float64]:
         data = self._tree.data
-        radec = xy_to_radec(np.transpose([data["x"], data["y"], data["z"]]))
+        radec = coord.xy_to_radec(np.transpose([data["x"], data["y"], data["z"]]))
 
         dtype = [("ra", "f8"), ("dec", "f8"), ("weight", "f8")]
         array = np.empty(len(data), dtype=dtype)
@@ -100,11 +45,11 @@ def sum_weight(self) -> float:
 
     @property
     def center(self) -> tuple(float, float, float):
-        return tuple(xy_to_radec(self._tree.center)[0])
+        return tuple(coord.xy_to_radec(self._tree.center)[0])
 
     @property
     def radius(self) -> float:
-        return radius_to_angle(self._tree.radius)
+        return coord.radius_to_angle(self._tree.radius)
 
     @classmethod
     def from_random(
@@ -115,11 +60,11 @@ def from_random(
         dec_max: float,
         size: int,
     ) -> AngularTree:
-        x_min, y_min = radec_to_xy(ra_min, dec_min)
-        x_max, y_max = radec_to_xy(ra_max, dec_max)
+        x_min, y_min = coord.radec_to_xy(ra_min, dec_min)
+        x_max, y_max = coord.radec_to_xy(ra_max, dec_max)
         x = np.random.uniform(x_min, x_max, size)
         y = np.random.uniform(y_min, y_max, size)
-        radec = xy_to_radec(np.transpose([x, y]))
+        radec = coord.xy_to_radec(np.transpose([x, y]))
         return cls(radec)
 
     @classmethod
@@ -139,32 +84,32 @@ def get_node_data(self) -> NDArray:
 
     def count_radius(
         self,
-        radec: NDArray[np.float64],
+        radec: NDArray,
         angle: float,
-        weight: NDArray[np.float64] | None = None,
+        weight: NDArray | None = None,
     ) -> float:
-        xyz = radec_to_xyz(radec)
-        radius = angle_to_radius(angle)
+        xyz = coord.radec_to_xyz(radec)
+        radius = coord.angle_to_radius(angle)
         return self._tree.count_radius(xyz, radius, weight)
 
     def count_range(
         self,
-        radec: NDArray[np.float64],
+        radec: NDArray,
         ang_min: float,
         ang_max: float,
-        weight: NDArray[np.float64] | None = None,
+        weight: NDArray | None = None,
     ) -> float:
-        xyz = radec_to_xyz(radec)
-        rmin = angle_to_radius(ang_min)
-        rmax = angle_to_radius(ang_max)
+        xyz = coord.radec_to_xyz(radec)
+        rmin = coord.angle_to_radius(ang_min)
+        rmax = coord.angle_to_radius(ang_max)
         return self._tree.count_range(xyz, rmin, rmax, weight)
 
     def dualcount_radius(
         self,
         other: AngularTree,
         angle: float,
     ) -> float:
-        radius = angle_to_radius(angle)
+        radius = coord.angle_to_radius(angle)
         return self._tree.dualcount_radius(other, radius)
 
     def dualcount_range(
@@ -173,6 +118,6 @@ def dualcount_range(
         ang_min: float,
         ang_max: float,
     ) -> float:
-        rmin = angle_to_radius(ang_min)
-        rmax = angle_to_radius(ang_max)
+        rmin = coord.angle_to_radius(ang_min)
+        rmax = coord.angle_to_radius(ang_max)
         return self._tree.dualcount_range(other, rmin, rmax)

balltree/coordinates.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+
+def sgn(val: ArrayLike) -> NDArray[np.float64]:
+    return np.where(val == 0, 1.0, np.sign(val, dtype=np.float64))
+
+
+def angular_to_cylinder(radec: NDArray) -> NDArray[np.float64]:
+    radec = np.atleast_2d(radec)
+    xy = np.empty_like(radec, dtype=np.float64)
+    xy[:, 0] = radec[:, 0]
+    xy[:, 1] = np.sin(radec[:, 1])
+    return xy
+
+
+def cylinder_to_angular(xy: NDArray) -> NDArray[np.float64]:
+    radec = np.empty_like(xy, dtype=np.float64)
+    radec[:, 0] = xy[:, 0]
+    radec[:, 1] = np.arcsin(xy[:, 1])
+    return radec
+
+
+def angular_to_euclidean(radec: NDArray) -> NDArray[np.float64]:
+    radec = np.atleast_2d(radec)
+    ra = radec[:, 0]
+    dec = radec[:, 1]
+    cos_dec = np.cos(dec)
+
+    xyz = np.empty((len(ra), 3), dtype=np.float64)
+    xyz[:, 0] = np.cos(ra) * cos_dec
+    xyz[:, 1] = np.sin(ra) * cos_dec
+    xyz[:, 2] = np.sin(dec)
+    return xyz
+
+
+def euclidean_to_angular(xyz: NDArray) -> NDArray[np.float64]:
+    xyz = np.atleast_2d(xyz)
+    x = xyz[:, 0]
+    y = xyz[:, 1]
+    z = xyz[:, 2]
+    r_d2 = np.sqrt(x * x + y * y)
+    r_d3 = np.sqrt(x * x + y * y + z * z)
+
+    radec = np.empty((len(x), 2), dtype=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)
+    return radec
+
+
+def angle_to_chorddist(angle: ArrayLike) -> NDArray[np.float64]:
+    return 2.0 * np.sin(angle / 2.0, dtype=np.float64)
+
+
+def chorddist_to_angle(chord: ArrayLike) -> NDArray[np.float64]:
+    return 2.0 * np.arcsin(chord / 2.0, dtype=np.float64)

pyproject.toml

@@ -30,7 +30,7 @@ classifiers = [
 ]
 requires-python = ">=3.7"
 dependencies = [
-    "numpy",
+    "numpy>=1.21",
 ]
 
 [project.urls]

requirements.txt

@@ -1,2 +1,2 @@
-numpy
+numpy>=1.21
 pytest

requirements_dev.txt

@@ -1,4 +1,4 @@
-numpy
+numpy>=1.21
 pytest
 black
 isort

tests/test_angulartree.py

@@ -0,0 +1,68 @@
+import pytest
+
+# from balltree.angulartree import AngularTree
+
+
+class TestAngularTree:
+    def test_init(self):
+        pass
+
+    @pytest.mark.skip
+    def test_data(self):
+        pass
+
+    @pytest.mark.skip
+    def test_num_data(self):
+        pass
+
+    @pytest.mark.skip
+    def test_leafsize(self):
+        pass
+
+    @pytest.mark.skip
+    def test_sum_weight(self):
+        pass
+
+    @pytest.mark.skip
+    def test_center(self):
+        pass
+
+    @pytest.mark.skip
+    def test_radius(self):
+        pass
+
+    @pytest.mark.skip
+    def test_from_random(self):
+        pass
+
+    @pytest.mark.skip
+    def test_from_file(self):
+        pass
+
+    @pytest.mark.skip
+    def test_to_file(self):
+        pass
+
+    @pytest.mark.skip
+    def test_count_nodes(self):
+        pass
+
+    @pytest.mark.skip
+    def test_get_node_data(self):
+        pass
+
+    @pytest.mark.skip
+    def test_count_radius(self):
+        pass
+
+    @pytest.mark.skip
+    def test_count_range(self):
+        pass
+
+    @pytest.mark.skip
+    def test_dualcount_radius(self):
+        pass
+
+    @pytest.mark.skip
+    def test_dualcount_range(self):
+        pass