Merge pull request #190 from astronomy-commons/issue/189/remove-slow-…

…wrapping-call

Add box filter benchmark on partition

benchmarks/benchmarks.py

@@ -6,9 +6,11 @@
 import os
 
 import numpy as np
+import pandas as pd
 
 import lsdb
 from benchmarks.utils import upsample_array
+from lsdb.core.search.box_search import box_filter
 from lsdb.core.search.polygon_search import get_cartesian_polygon
 
 TEST_DIR = os.path.join(os.path.dirname(__file__), "..", "tests")
@@ -51,3 +53,12 @@ def time_polygon_search():
     polygon, _ = get_cartesian_polygon(vertices)
     # Apply vectorized filtering on the catalog points
     polygon.contains(np.radians(catalog_ra), np.radians(catalog_dec))
+
+
+def time_box_filter_on_partition():
+    """Time box search on a single partition"""
+    metadata = load_small_sky_order1().hc_structure
+    mock_partition_df = pd.DataFrame(
+        np.linspace(-1000, 1000, 100_000), columns=[metadata.catalog_info.ra_column]
+    )
+    box_filter(mock_partition_df, ra=(-20, 40), dec=None, metadata=metadata).compute()

src/lsdb/core/search/box_search.py

@@ -7,7 +7,7 @@
 import numpy as np
 import pandas as pd
 from hipscat.pixel_math import HealpixPixel
-from hipscat.pixel_math.box_filter import filter_pixels_by_box, wrap_angles, wrap_ra_values
+from hipscat.pixel_math.box_filter import filter_pixels_by_box, wrap_ra_angles
 from hipscat.pixel_math.validators import validate_box_search
 from hipscat.pixel_tree.pixel_tree_builder import PixelTreeBuilder
 
@@ -28,7 +28,7 @@ def __init__(
         ra: Tuple[float, float] | None = None,
         dec: Tuple[float, float] | None = None,
     ):
-        ra = wrap_ra_values(ra)
+        ra = tuple(wrap_ra_angles(ra)) if ra else None
         validate_box_search(ra, dec)
 
         self.ra, self.dec = ra, dec
@@ -65,7 +65,7 @@ def box_filter(
     mask = np.ones(len(data_frame), dtype=bool)
     if ra is not None:
         ra_values = data_frame[metadata.catalog_info.ra_column]
-        wrapped_ra = np.asarray(wrap_angles(ra_values))
+        wrapped_ra = np.asarray(wrap_ra_angles(ra_values))
         mask_ra = _create_ra_mask(ra, wrapped_ra)
         mask = np.logical_and(mask, mask_ra)
     if dec is not None:

src/lsdb/core/search/polygon_search.py

@@ -67,7 +67,7 @@ def get_cartesian_polygon(
     """Creates the convex polygon to filter pixels with. It transforms the vertices, provided
     in sky coordinates of ra and dec, to their respective cartesian representation on the unit sphere.
 
-    Arguments:
+    Args:
         vertices (List[Tuple[float, float]): The list of vertices of the polygon to
             filter pixels with, as a list of (ra,dec) coordinates, in degrees.