Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use numba jit compilation instead of precompilation #395

Merged
merged 8 commits into from
Oct 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ build-backend = "setuptools.build_meta"
write_to = "src/hats/_version.py"

[tool.pytest.ini_options]
timeout = 1
timeout = 5
testpaths = [
"tests",
]
Expand Down
2 changes: 1 addition & 1 deletion src/hats/pixel_math/margin_bounding.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ def check_margin_bounds(r_asc, dec, pixel_order, pixel, margin_threshold, step=1


# numba jit compiler doesn't count for coverage tests, so we'll set no cover.
@njit("double(double[:], double[:], double)")
@njit
def _find_minimum_distance(separations, distances, margin_threshold): # pragma: no cover
"""Find the minimum distance between a given datapoint and a healpixel"""
minimum_index = np.argmin(separations)
Expand Down
8 changes: 1 addition & 7 deletions src/hats/pixel_tree/moc_filter.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import numba
import numpy as np
from mocpy import MOC
from numba import njit
Expand Down Expand Up @@ -26,12 +25,7 @@ def filter_by_moc(
return PixelTree(tree.tree[tree_mask], tree.tree_order)


@njit(
numba.bool_[::1](
numba.int64[:, :],
numba.uint64[:, :],
)
)
@njit
def perform_filter_by_moc(
tree: np.ndarray,
moc: np.ndarray,
Expand Down
25 changes: 4 additions & 21 deletions src/hats/pixel_tree/pixel_alignment.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

from typing import Callable, Dict, List

import numba
import numpy as np
import pandas as pd
from mocpy import MOC
Expand Down Expand Up @@ -141,7 +140,7 @@ def get_pixel_mapping_df(mapping: np.ndarray, map_order: int) -> pd.DataFrame:


# pylint: disable=too-many-statements
@njit(numba.int64[::1, :](numba.int64[:, :], numba.int64[:, :]))
@njit
def perform_inner_align_trees(
left: np.ndarray,
right: np.ndarray,
Expand Down Expand Up @@ -208,15 +207,7 @@ def perform_inner_align_trees(
return mapping[:out_index].T


@njit(
numba.types.void(
numba.int64,
numba.int64,
numba.int64[:],
numba.boolean,
numba.types.List(numba.int64[::1]),
)
)
@njit
def _add_pixels_until(
add_from: int,
add_to: int,
Expand Down Expand Up @@ -259,15 +250,7 @@ def _add_pixels_until(
add_from = add_from + pixel_size


@njit(
numba.types.void(
numba.int64,
numba.int64[:, :],
numba.int64,
numba.boolean,
numba.types.List(numba.int64[::1]),
)
)
@njit
def _add_remaining_pixels(
added_until: int,
pixel_list: np.ndarray,
Expand Down Expand Up @@ -303,7 +286,7 @@ def _add_remaining_pixels(


# pylint: disable=too-many-statements
@njit(numba.types.List(numba.int64[::1])(numba.int64[:, :], numba.int64[:, :], numba.boolean, numba.boolean))
@njit
def perform_align_trees(
left: np.ndarray,
right: np.ndarray,
Expand Down
3 changes: 2 additions & 1 deletion tests/hats/catalog/test_catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,6 @@ def test_cone_filter_big(small_sky_order1_catalog):
assert (1, 47) in filtered_catalog.pixel_tree


@pytest.mark.timeout(5)
def test_cone_filter_multiple_order(catalog_info):
catalog_pixel_list = [
HealpixPixel(6, 30),
Expand Down Expand Up @@ -495,6 +494,7 @@ def test_empty_directory(tmp_path, catalog_info_data):
assert catalog.catalog_name == "test_name"


@pytest.mark.timeout(20)
def test_generate_negative_tree_pixels(small_sky_order1_catalog):
"""Test generate_negative_tree_pixels on a basic catalog."""
expected_pixels = [
Expand All @@ -516,6 +516,7 @@ def test_generate_negative_tree_pixels(small_sky_order1_catalog):
assert negative_tree == expected_pixels


@pytest.mark.timeout(20)
def test_generate_negative_tree_pixels_order_0(small_sky_catalog):
"""Test generate_negative_tree_pixels on a catalog with only order 0 pixels."""
expected_pixels = [
Expand Down
2 changes: 2 additions & 0 deletions tests/hats/pixel_tree/test_pixel_alignment.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import numpy as np
import pytest

from hats.catalog import Catalog
from hats.pixel_math import HealpixPixel
Expand Down Expand Up @@ -269,6 +270,7 @@ def test_catalog_align_outer(pixel_tree_2, pixel_tree_3, aligned_trees_2_3_outer
assert alignment.moc == moc


@pytest.mark.timeout(20)
def test_outer_align_start_0():
left_tree = PixelTree.from_healpix([HealpixPixel(0, 0)])
right_tree = PixelTree.from_healpix([HealpixPixel(1, 1)])
Expand Down