diff --git a/ranx/__init__.py b/ranx/__init__.py index ad2ce56..a2d0749 100644 --- a/ranx/__init__.py +++ b/ranx/__init__.py @@ -7,12 +7,21 @@ "plot", "Qrels", "Run", + "use_numba", + "set_numba_enabled", ] -from numba import config - +from .config import set_numba_enabled, use_numba from .data_structures import Qrels, Run from .meta import compare, evaluate, fuse, normalize, optimize_fusion, plot -# Set numba threading layer to workqueue -config.THREADING_LAYER = "workqueue" +# Conditional Numba configuration +if use_numba(): + try: + from numba import config + + # Set numba threading layer to workqueue + config.THREADING_LAYER = "workqueue" + except ImportError: + # Numba not available, silently continue + pass diff --git a/ranx/config.py b/ranx/config.py new file mode 100644 index 0000000..5ef81bd --- /dev/null +++ b/ranx/config.py @@ -0,0 +1,46 @@ +"""Configuration system for ranx library.""" + +import os +from typing import Optional + +# Global configuration state +_USE_NUMBA: Optional[bool] = None + + +def use_numba() -> bool: + """ + Check if Numba should be used for performance optimizations. + + Returns True by default, but can be disabled via: + 1. Environment variable: RANX_USE_NUMBA=false + 2. Programmatically: set_numba_enabled(False) + + Returns: + bool: True if Numba should be used, False otherwise + """ + global _USE_NUMBA + if _USE_NUMBA is None: + env_value = os.environ.get("RANX_USE_NUMBA", "true").lower() + _USE_NUMBA = env_value not in ("false", "0", "no", "off") + return _USE_NUMBA + + +def set_numba_enabled(enabled: bool) -> None: + """ + Programmatically enable or disable Numba usage. + + Args: + enabled: True to enable Numba, False to disable + """ + global _USE_NUMBA + _USE_NUMBA = enabled + + +def reset_numba_config() -> None: + """ + Reset Numba configuration to default (reads from environment again). + + Primarily used for testing purposes. + """ + global _USE_NUMBA + _USE_NUMBA = None diff --git a/ranx/data_structures/generic.py b/ranx/data_structures/generic.py index 115b0fc..cc91c11 100644 --- a/ranx/data_structures/generic.py +++ b/ranx/data_structures/generic.py @@ -1,24 +1,29 @@ -from numba import njit, types -from numba.typed import Dict as TypedDict -from numba.typed import List as TypedList +from ..decorators import create_typed_dict, create_typed_list + +# Handle Numba-specific imports conditionally +try: + from numba import types + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False -@njit(cache=True) def create_empty_results_dict(): - return TypedDict.empty( - key_type=types.unicode_type, - value_type=types.float64, + return create_typed_dict( + key_type=types.unicode_type if NUMBA_AVAILABLE else None, + value_type=types.float64 if NUMBA_AVAILABLE else None, ) -@njit(cache=True) def create_empty_results_dict_list(length): - return TypedList([create_empty_results_dict() for _ in range(length)]) + return create_typed_list( + initial_list=[create_empty_results_dict() for _ in range(length)] + ) -@njit(cache=True) def convert_results_dict_list_to_run(q_ids, results_dict_list): - combined_run = TypedDict() + combined_run = create_typed_dict() for i, q_id in enumerate(q_ids): combined_run[q_id] = results_dict_list[i] diff --git a/ranx/data_structures/qrels.py b/ranx/data_structures/qrels.py index f0d24bf..5ee25e5 100644 --- a/ranx/data_structures/qrels.py +++ b/ranx/data_structures/qrels.py @@ -7,9 +7,21 @@ import numpy as np import orjson import pandas as pd -from numba import njit, prange, types -from numba.typed import Dict as TypedDict -from numba.typed import List as TypedList + +try: + from numba import njit, prange, types + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + + +try: + from numba.typed import Dict as TypedDict + from numba.typed import List as TypedList +except ImportError: + TypedDict = dict + TypedList = list from .common import ( add_and_sort, diff --git a/ranx/data_structures/run.py b/ranx/data_structures/run.py index 381ea81..c3f794f 100644 --- a/ranx/data_structures/run.py +++ b/ranx/data_structures/run.py @@ -5,9 +5,17 @@ import numpy as np import pandas as pd -from numba import types -from numba.typed import Dict as TypedDict -from numba.typed import List as TypedList + +try: + from numba import types + from numba.typed import Dict as TypedDict + from numba.typed import List as TypedList + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + TypedDict = dict + TypedList = list from ..io import download, load_json, load_lz4, save_json, save_lz4 from .common import ( diff --git a/ranx/decorators.py b/ranx/decorators.py new file mode 100644 index 0000000..8ff023d --- /dev/null +++ b/ranx/decorators.py @@ -0,0 +1,135 @@ +"""Conditional decorators for optional Numba support.""" + +from .config import use_numba + +# Check if Numba is available +try: + from numba import jit, njit + from numba.typed import Dict as TypedDict + from numba.typed import List as TypedList + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + # Create dummy objects for when Numba is not available + TypedDict = dict + TypedList = list + + +def maybe_njit(*args, **kwargs): + """ + Conditional njit decorator that falls back to identity function when Numba is disabled. + + This decorator will apply Numba's njit compilation when: + 1. Numba is available (installed) + 2. Numba is enabled via configuration + + Otherwise, it returns the function unchanged (pure Python execution). + + Args: + *args, **kwargs: Same arguments as numba.njit + + Returns: + Function decorator that conditionally applies Numba compilation + """ + + def decorator(func): + if NUMBA_AVAILABLE and use_numba(): + return njit(*args, **kwargs)(func) + else: + return func + + # Handle the case where maybe_njit is used without arguments: @maybe_njit + if len(args) == 1 and callable(args[0]) and not kwargs: + func = args[0] + if NUMBA_AVAILABLE and use_numba(): + return njit()(func) + else: + return func + + return decorator + + +def maybe_jit(*args, **kwargs): + """ + Conditional jit decorator that falls back to identity function when Numba is disabled. + + Similar to maybe_njit but uses jit instead of njit. + + Args: + *args, **kwargs: Same arguments as numba.jit + + Returns: + Function decorator that conditionally applies Numba compilation + """ + + def decorator(func): + if NUMBA_AVAILABLE and use_numba(): + return jit(*args, **kwargs)(func) + else: + return func + + # Handle the case where maybe_jit is used without arguments: @maybe_jit + if len(args) == 1 and callable(args[0]) and not kwargs: + func = args[0] + if NUMBA_AVAILABLE and use_numba(): + return jit()(func) + else: + return func + + return decorator + + +# Note: prange requires separate implementations because Numba needs compile-time +# knowledge of whether to use prange or range. See metrics files for examples. + + +def create_typed_dict(key_type=None, value_type=None, initial_dict=None): + """ + Create a typed dictionary that falls back to regular dict when Numba is disabled. + + Args: + key_type: Numba type for keys (ignored when Numba disabled) + value_type: Numba type for values (ignored when Numba disabled) + initial_dict: Initial dictionary to populate + + Returns: + numba.typed.Dict when Numba enabled, regular dict otherwise + """ + if NUMBA_AVAILABLE and use_numba(): + if initial_dict: + typed_dict = TypedDict() + for k, v in initial_dict.items(): + typed_dict[k] = v + return typed_dict + elif key_type is not None and value_type is not None: + return TypedDict.empty(key_type, value_type) + else: + return TypedDict() + else: + return dict(initial_dict) if initial_dict else {} + + +def create_typed_list(item_type=None, initial_list=None): + """ + Create a typed list that falls back to regular list when Numba is disabled. + + Args: + item_type: Numba type for items (ignored when Numba disabled) + initial_list: Initial list to populate + + Returns: + numba.typed.List when Numba enabled, regular list otherwise + """ + if NUMBA_AVAILABLE and use_numba(): + if initial_list: + typed_list = TypedList() + for item in initial_list: + typed_list.append(item) + return typed_list + elif item_type is not None: + return TypedList.empty_list(item_type) + else: + return TypedList() + else: + return list(initial_list) if initial_list else [] diff --git a/ranx/metrics/average_precision.py b/ranx/metrics/average_precision.py index 5c4e10f..3c439d4 100644 --- a/ranx/metrics/average_precision.py +++ b/ranx/metrics/average_precision.py @@ -1,14 +1,13 @@ from typing import Union -import numba import numpy as np -from numba import njit, prange +from ..decorators import maybe_njit from .common import clean_qrels, fix_k # LOW LEVEL FUNCTIONS ========================================================== -@njit(cache=True) +@maybe_njit(cache=True) def _average_precision(qrels, run, k, rel_lvl): qrels = clean_qrels(qrels, rel_lvl) if len(qrels) == 0: @@ -42,18 +41,44 @@ def _average_precision(qrels, run, k, rel_lvl): return np.sum(precision_scores) / qrels.shape[0] -@njit(cache=True, parallel=True) -def _average_precision_parallel(qrels, run, k, rel_lvl): +# Handle parallel version with conditional compilation +try: + from numba import njit, prange + + @njit(cache=True, parallel=True) + def _average_precision_parallel_numba(qrels, run, k, rel_lvl): + scores = np.zeros((len(qrels)), dtype=np.float64) + for i in prange(len(qrels)): + scores[i] = _average_precision(qrels[i], run[i], k, rel_lvl) + return scores + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + + +def _average_precision_numpy(qrels, run, k, rel_lvl): + """NumPy fallback implementation.""" scores = np.zeros((len(qrels)), dtype=np.float64) - for i in prange(len(qrels)): + for i in range(len(qrels)): scores[i] = _average_precision(qrels[i], run[i], k, rel_lvl) return scores +def _average_precision_parallel(qrels, run, k, rel_lvl): + """Dispatch to best available implementation.""" + from ..config import use_numba + + if NUMBA_AVAILABLE and use_numba(): + return _average_precision_parallel_numba(qrels, run, k, rel_lvl) + else: + return _average_precision_numpy(qrels, run, k, rel_lvl) + + # HIGH LEVEL FUNCTIONS ========================================================= def average_precision( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -72,9 +97,9 @@ def average_precision( - $R$ is the total number of relevant documents. Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels: IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run: IDs and relevance scores of _retrieved_ documents. k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0. diff --git a/ranx/metrics/common.py b/ranx/metrics/common.py index 763e3f0..f0a4505 100644 --- a/ranx/metrics/common.py +++ b/ranx/metrics/common.py @@ -1,12 +1,13 @@ import numpy as np -from numba import njit +from ..decorators import maybe_njit -@njit(cache=True) + +@maybe_njit(cache=True) def clean_qrels(qrels, rel_lvl): return qrels[np.argwhere(qrels[:, 1] >= rel_lvl).flatten()] -@njit(cache=True) +@maybe_njit(cache=True) def fix_k(k, run): return run.shape[0] if k == 0 or k > run.shape[0] else k diff --git a/ranx/metrics/f1.py b/ranx/metrics/f1.py index 2a4d7aa..6483480 100644 --- a/ranx/metrics/f1.py +++ b/ranx/metrics/f1.py @@ -1,15 +1,14 @@ from typing import Union -import numba import numpy as np -from numba import njit, prange +from ..decorators import maybe_njit from .common import clean_qrels from .hits import _hits # LOW LEVEL FUNCTIONS ========================================================== -@njit(cache=True) +@maybe_njit(cache=True) def _f1(qrels, run, k, rel_lvl): qrels = clean_qrels(qrels, rel_lvl) if len(qrels) == 0: @@ -29,18 +28,44 @@ def _f1(qrels, run, k, rel_lvl): return 2 * ((precision_score * recall_score) / (precision_score + recall_score)) -@njit(cache=True, parallel=True) -def _f1_parallel(qrels, run, k, rel_lvl): +# Handle parallel version with conditional compilation +try: + from numba import njit, prange + + @njit(cache=True, parallel=True) + def _f1_parallel_numba(qrels, run, k, rel_lvl): + scores = np.zeros((len(qrels)), dtype=np.float64) + for i in prange(len(qrels)): + scores[i] = _f1(qrels[i], run[i], k, rel_lvl) + return scores + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + + +def _f1_numpy(qrels, run, k, rel_lvl): + """NumPy fallback implementation.""" scores = np.zeros((len(qrels)), dtype=np.float64) - for i in prange(len(qrels)): + for i in range(len(qrels)): scores[i] = _f1(qrels[i], run[i], k, rel_lvl) return scores +def _f1_parallel(qrels, run, k, rel_lvl): + """Dispatch to best available implementation.""" + from ..config import use_numba + + if NUMBA_AVAILABLE and use_numba(): + return _f1_parallel_numba(qrels, run, k, rel_lvl) + else: + return _f1_numpy(qrels, run, k, rel_lvl) + + # HIGH LEVEL FUNCTIONS ========================================================= def f1( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -63,9 +88,9 @@ def f1( $$ Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels: IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run: IDs and relevance scores of _retrieved_ documents. k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0. diff --git a/ranx/metrics/hits.py b/ranx/metrics/hits.py index 26895f8..ef4971c 100644 --- a/ranx/metrics/hits.py +++ b/ranx/metrics/hits.py @@ -1,14 +1,13 @@ from typing import Union -import numba import numpy as np -from numba import njit, prange +from ..decorators import maybe_njit from .common import clean_qrels, fix_k # LOW LEVEL FUNCTIONS ========================================================== -@njit(cache=True) +@maybe_njit(cache=True) def _hits(qrels, run, k, rel_lvl): qrels = clean_qrels(qrels, rel_lvl) if len(qrels) == 0: @@ -34,18 +33,44 @@ def _hits(qrels, run, k, rel_lvl): return hits -@njit(cache=True, parallel=True) -def _hits_parallel(qrels, run, k, rel_lvl): +# Handle parallel version with conditional compilation +try: + from numba import njit, prange + + @njit(cache=True, parallel=True) + def _hits_parallel_numba(qrels, run, k, rel_lvl): + scores = np.zeros((len(qrels)), dtype=np.float64) + for i in prange(len(qrels)): + scores[i] = _hits(qrels[i], run[i], k, rel_lvl) + return scores + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + + +def _hits_numpy(qrels, run, k, rel_lvl): + """NumPy fallback implementation.""" scores = np.zeros((len(qrels)), dtype=np.float64) - for i in prange(len(qrels)): + for i in range(len(qrels)): scores[i] = _hits(qrels[i], run[i], k, rel_lvl) return scores +def _hits_parallel(qrels, run, k, rel_lvl): + """Dispatch to best available implementation.""" + from ..config import use_numba + + if NUMBA_AVAILABLE and use_numba(): + return _hits_parallel_numba(qrels, run, k, rel_lvl) + else: + return _hits_numpy(qrels, run, k, rel_lvl) + + # HIGH LEVEL FUNCTIONS ========================================================= def hits( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -55,9 +80,9 @@ def hits( If k > 0, only the top-k retrieved documents are considered. Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels: IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run: IDs and relevance scores of _retrieved_ documents. k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0. diff --git a/ranx/metrics/ndcg.py b/ranx/metrics/ndcg.py index c51b109..dc789da 100644 --- a/ranx/metrics/ndcg.py +++ b/ranx/metrics/ndcg.py @@ -1,14 +1,13 @@ from typing import Union -import numba import numpy as np -from numba import njit, prange +from ..decorators import maybe_njit from .common import clean_qrels, fix_k # LOW LEVEL FUNCTIONS ========================================================== -@njit(cache=True) +@maybe_njit(cache=True) def _dcg(qrels, run, k, rel_lvl, jarvelin): qrels = clean_qrels(qrels, rel_lvl) if len(qrels) == 0: @@ -40,20 +39,53 @@ def _dcg(qrels, run, k, rel_lvl, jarvelin): return np.sum((2**weighted_hit_list - 1) / np.log2(np.arange(1, k + 1) + 1)) -@njit(cache=True, parallel=True) -def _dcg_parallel(qrels, run, k, rel_lvl, jarvelin): +# Handle parallel version with conditional compilation +try: + from numba import njit, prange + + @njit(cache=True, parallel=True) + def _dcg_parallel_numba(qrels, run, k, rel_lvl, jarvelin): + scores = np.zeros((len(qrels)), dtype=np.float64) + for i in prange(len(qrels)): + scores[i] = _dcg(qrels[i], run[i], k, rel_lvl, jarvelin) + return scores + + @njit(cache=True, parallel=True) + def _ndcg_parallel_numba(qrels, run, k, rel_lvl, jarvelin): + scores = np.zeros((len(qrels)), dtype=np.float64) + for i in prange(len(qrels)): + scores[i] = _ndcg(qrels[i], run[i], k, rel_lvl, jarvelin) + return scores + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + + +def _dcg_numpy(qrels, run, k, rel_lvl, jarvelin): + """NumPy fallback implementation.""" scores = np.zeros((len(qrels)), dtype=np.float64) - for i in prange(len(qrels)): + for i in range(len(qrels)): scores[i] = _dcg(qrels[i], run[i], k, rel_lvl, jarvelin) return scores -@njit(cache=True) +def _dcg_parallel(qrels, run, k, rel_lvl, jarvelin): + """Dispatch to best available implementation.""" + from ..config import use_numba + + if NUMBA_AVAILABLE and use_numba(): + return _dcg_parallel_numba(qrels, run, k, rel_lvl, jarvelin) + else: + return _dcg_numpy(qrels, run, k, rel_lvl, jarvelin) + + +@maybe_njit(cache=True) def _idcg(qrels, k, rel_lvl, jarvelin): return _dcg(qrels, qrels, k, rel_lvl, jarvelin) -@njit(cache=True) +@maybe_njit(cache=True) def _ndcg(qrels, run, k, rel_lvl, jarvelin): dcg_score = _dcg(qrels, run, k, rel_lvl, jarvelin) idcg_score = _idcg(qrels, k, rel_lvl, jarvelin) @@ -65,18 +97,28 @@ def _ndcg(qrels, run, k, rel_lvl, jarvelin): return dcg_score / idcg_score -@njit(cache=True, parallel=True) -def _ndcg_parallel(qrels, run, k, rel_lvl, jarvelin): +def _ndcg_numpy(qrels, run, k, rel_lvl, jarvelin): + """NumPy fallback implementation.""" scores = np.zeros((len(qrels)), dtype=np.float64) - for i in prange(len(qrels)): + for i in range(len(qrels)): scores[i] = _ndcg(qrels[i], run[i], k, rel_lvl, jarvelin) return scores +def _ndcg_parallel(qrels, run, k, rel_lvl, jarvelin): + """Dispatch to best available implementation.""" + from ..config import use_numba + + if NUMBA_AVAILABLE and use_numba(): + return _ndcg_parallel_numba(qrels, run, k, rel_lvl, jarvelin) + else: + return _ndcg_numpy(qrels, run, k, rel_lvl, jarvelin) + + # HIGH LEVEL FUNCTIONS ========================================================= def dcg( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -105,9 +147,9 @@ def dcg( ``` Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels: IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run: IDs and relevance scores of _retrieved_ documents. k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0. @@ -124,8 +166,8 @@ def dcg( def ndcg( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -168,9 +210,9 @@ def ndcg( ``` Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels: IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run: IDs and relevance scores of _retrieved_ documents. k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0. @@ -187,8 +229,8 @@ def ndcg( def dcg_burges( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -223,9 +265,9 @@ def dcg_burges( ``` Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels: IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run: IDs and relevance scores of _retrieved_ documents. k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0. @@ -242,8 +284,8 @@ def dcg_burges( def ndcg_burges( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -292,9 +334,9 @@ def ndcg_burges( ``` Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels: IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run: IDs and relevance scores of _retrieved_ documents. k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0. diff --git a/ranx/metrics/precision.py b/ranx/metrics/precision.py index fb50435..b78c997 100644 --- a/ranx/metrics/precision.py +++ b/ranx/metrics/precision.py @@ -1,14 +1,13 @@ from typing import Union -import numba import numpy as np -from numba import njit, prange +from ..decorators import maybe_njit from .hits import _hits # LOW LEVEL FUNCTIONS ========================================================== -@njit(cache=True) +@maybe_njit(cache=True) def _precision(qrels, run, k, rel_lvl): k = k if k != 0 else run.shape[0] if k == 0: @@ -17,18 +16,44 @@ def _precision(qrels, run, k, rel_lvl): return _hits(qrels, run, k, rel_lvl) / k -@njit(cache=True, parallel=True) -def _precision_parallel(qrels, run, k, rel_lvl): +# Handle parallel version with conditional compilation +try: + from numba import njit, prange + + @njit(cache=True, parallel=True) + def _precision_parallel_numba(qrels, run, k, rel_lvl): + scores = np.zeros((len(qrels)), dtype=np.float64) + for i in prange(len(qrels)): + scores[i] = _precision(qrels[i], run[i], k, rel_lvl) + return scores + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + + +def _precision_numpy(qrels, run, k, rel_lvl): + """NumPy fallback implementation.""" scores = np.zeros((len(qrels)), dtype=np.float64) - for i in prange(len(qrels)): + for i in range(len(qrels)): scores[i] = _precision(qrels[i], run[i], k, rel_lvl) return scores +def _precision_parallel(qrels, run, k, rel_lvl): + """Dispatch to best available implementation.""" + from ..config import use_numba + + if NUMBA_AVAILABLE and use_numba(): + return _precision_parallel_numba(qrels, run, k, rel_lvl) + else: + return _precision_numpy(qrels, run, k, rel_lvl) + + # HIGH LEVEL FUNCTIONS ========================================================= def precision( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -59,9 +84,9 @@ def precision( - $r_k$ is the number of retrieved relevant documents at k. Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels (Union[np.ndarray, list]): IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run (Union[np.ndarray, list]): IDs and relevance scores of _retrieved_ documents. k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0. diff --git a/ranx/metrics/recall.py b/ranx/metrics/recall.py index 618d07e..37a2c39 100644 --- a/ranx/metrics/recall.py +++ b/ranx/metrics/recall.py @@ -1,15 +1,14 @@ from typing import Union -import numba import numpy as np -from numba import njit, prange +from ..decorators import maybe_njit from .common import clean_qrels from .hits import _hits # LOW LEVEL FUNCTIONS ========================================================== -@njit(cache=True) +@maybe_njit(cache=True) def _recall(qrels, run, k, rel_lvl): qrels = clean_qrels(qrels, rel_lvl) if len(qrels) == 0: @@ -22,18 +21,44 @@ def _recall(qrels, run, k, rel_lvl): return _hits(qrels, run, k, rel_lvl) / qrels.shape[0] -@njit(cache=True, parallel=True) -def _recall_parallel(qrels, run, k, rel_lvl): +# Handle parallel version with conditional compilation +try: + from numba import njit, prange + + @njit(cache=True, parallel=True) + def _recall_parallel_numba(qrels, run, k, rel_lvl): + scores = np.zeros((len(qrels)), dtype=np.float64) + for i in prange(len(qrels)): + scores[i] = _recall(qrels[i], run[i], k, rel_lvl) + return scores + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + + +def _recall_numpy(qrels, run, k, rel_lvl): + """NumPy fallback implementation.""" scores = np.zeros((len(qrels)), dtype=np.float64) - for i in prange(len(qrels)): + for i in range(len(qrels)): scores[i] = _recall(qrels[i], run[i], k, rel_lvl) return scores +def _recall_parallel(qrels, run, k, rel_lvl): + """Dispatch to best available implementation.""" + from ..config import use_numba + + if NUMBA_AVAILABLE and use_numba(): + return _recall_parallel_numba(qrels, run, k, rel_lvl) + else: + return _recall_numpy(qrels, run, k, rel_lvl) + + # HIGH LEVEL FUNCTIONS ========================================================= def recall( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -65,9 +90,9 @@ def recall( - $R$ is the total number of relevant documents. Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels: IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run: IDs and relevance scores of _retrieved_ documents. k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0. diff --git a/ranx/metrics/reciprocal_rank.py b/ranx/metrics/reciprocal_rank.py index 45872d0..8bf50ac 100644 --- a/ranx/metrics/reciprocal_rank.py +++ b/ranx/metrics/reciprocal_rank.py @@ -1,14 +1,13 @@ from typing import Union -import numba import numpy as np -from numba import njit, prange +from ..decorators import maybe_njit from .common import clean_qrels, fix_k # LOW LEVEL FUNCTIONS ========================================================== -@njit(cache=True) +@maybe_njit(cache=True) def _reciprocal_rank(qrels, run, k, rel_lvl): qrels = clean_qrels(qrels, rel_lvl) if len(qrels) == 0: @@ -22,18 +21,44 @@ def _reciprocal_rank(qrels, run, k, rel_lvl): return 0.0 -@njit(cache=True, parallel=True) -def _reciprocal_rank_parallel(qrels, run, k, rel_lvl): +# Handle parallel version with conditional compilation +try: + from numba import njit, prange + + @njit(cache=True, parallel=True) + def _reciprocal_rank_parallel_numba(qrels, run, k, rel_lvl): + scores = np.zeros((len(qrels)), dtype=np.float64) + for i in prange(len(qrels)): + scores[i] = _reciprocal_rank(qrels[i], run[i], k, rel_lvl) + return scores + + NUMBA_AVAILABLE = True +except ImportError: + NUMBA_AVAILABLE = False + + +def _reciprocal_rank_numpy(qrels, run, k, rel_lvl): + """NumPy fallback implementation.""" scores = np.zeros((len(qrels)), dtype=np.float64) - for i in prange(len(qrels)): + for i in range(len(qrels)): scores[i] = _reciprocal_rank(qrels[i], run[i], k, rel_lvl) return scores +def _reciprocal_rank_parallel(qrels, run, k, rel_lvl): + """Dispatch to best available implementation.""" + from ..config import use_numba + + if NUMBA_AVAILABLE and use_numba(): + return _reciprocal_rank_parallel_numba(qrels, run, k, rel_lvl) + else: + return _reciprocal_rank_numpy(qrels, run, k, rel_lvl) + + # HIGH LEVEL FUNCTIONS ========================================================= def reciprocal_rank( - qrels: Union[np.ndarray, numba.typed.List], - run: Union[np.ndarray, numba.typed.List], + qrels: Union[np.ndarray, list], + run: Union[np.ndarray, list], k: int = 0, rel_lvl: int = 1, ) -> np.ndarray: @@ -51,9 +76,9 @@ def reciprocal_rank( - $rank$ is the position of the first retrieved relevant document. Args: - qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents. + qrels: IDs and relevance scores of _relevant_ documents. - run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents. + run: IDs and relevance scores of _retrieved_ documents. k (int, optional): This argument is ignored. It was added to standardize metrics' input. Defaults to 0. diff --git a/tests/unit/ranx/test_numba_optional.py b/tests/unit/ranx/test_numba_optional.py new file mode 100644 index 0000000..2da385c --- /dev/null +++ b/tests/unit/ranx/test_numba_optional.py @@ -0,0 +1,128 @@ +"""Tests for optional Numba functionality.""" + +import pytest + +import ranx +from ranx.config import reset_numba_config + + +class TestNumbaOptional: + """Test that ranx works with Numba both enabled and disabled.""" + + def setup_method(self): + """Reset configuration before each test.""" + reset_numba_config() + + def test_precision_with_numba_enabled(self): + """Test precision calculation with Numba enabled.""" + ranx.set_numba_enabled(True) + + qrels_dict = {"q1": {"d1": 1, "d2": 1, "d3": 0}} + run_dict = {"q1": {"d1": 0.9, "d2": 0.8, "d3": 0.7}} + + qrels = ranx.Qrels.from_dict(qrels_dict) + run = ranx.Run.from_dict(run_dict) + + result = ranx.evaluate(qrels, run, ["precision@2"]) + + # Should get 2/2 = 1.0 precision (both top-2 docs are relevant) + assert result == pytest.approx(1.0) + + def test_precision_with_numba_disabled(self): + """Test precision calculation with Numba disabled.""" + ranx.set_numba_enabled(False) + + qrels_dict = {"q1": {"d1": 1, "d2": 1, "d3": 0}} + run_dict = {"q1": {"d1": 0.9, "d2": 0.8, "d3": 0.7}} + + qrels = ranx.Qrels.from_dict(qrels_dict) + run = ranx.Run.from_dict(run_dict) + + result = ranx.evaluate(qrels, run, ["precision@2"]) + + # Should get the same result as with Numba enabled + assert result == pytest.approx(1.0) + + def test_recall_with_numba_disabled(self): + """Test recall calculation with Numba disabled.""" + ranx.set_numba_enabled(False) + + qrels_dict = {"q1": {"d1": 1, "d2": 1, "d3": 1}} # 3 relevant docs + run_dict = {"q1": {"d1": 0.9, "d2": 0.8, "d3": 0.7}} + + qrels = ranx.Qrels.from_dict(qrels_dict) + run = ranx.Run.from_dict(run_dict) + + result = ranx.evaluate(qrels, run, ["recall@2"]) + + # Should get 2/3 ≈ 0.667 recall (found 2 out of 3 relevant docs in top-2) + assert result == pytest.approx(2 / 3) + + def test_hits_with_numba_disabled(self): + """Test hits calculation with Numba disabled.""" + ranx.set_numba_enabled(False) + + qrels_dict = {"q1": {"d1": 1, "d2": 1, "d3": 0}} + run_dict = {"q1": {"d1": 0.9, "d2": 0.8, "d3": 0.7}} + + qrels = ranx.Qrels.from_dict(qrels_dict) + run = ranx.Run.from_dict(run_dict) + + result = ranx.evaluate(qrels, run, ["hits@2"]) + + # Should get 2 hits (both top-2 docs are relevant) + assert result == pytest.approx(2.0) + + def test_multiple_queries_with_numba_disabled(self): + """Test multiple queries with Numba disabled.""" + ranx.set_numba_enabled(False) + + qrels_dict = {"q1": {"d1": 1, "d2": 1}, "q2": {"d3": 1, "d4": 1, "d5": 1}} + run_dict = { + "q1": {"d1": 0.9, "d2": 0.8, "d6": 0.7}, + "q2": {"d3": 0.95, "d4": 0.85, "d5": 0.75}, + } + + qrels = ranx.Qrels.from_dict(qrels_dict) + run = ranx.Run.from_dict(run_dict) + + result = ranx.evaluate(qrels, run, ["precision@2", "recall@2"]) + + # q1: precision@2 = 2/2 = 1.0, recall@2 = 2/2 = 1.0 + # q2: precision@2 = 2/2 = 1.0, recall@2 = 2/3 = 0.667 + # Mean precision@2 = (1.0 + 1.0) / 2 = 1.0 + # Mean recall@2 = (1.0 + 0.667) / 2 = 0.833 + assert result["precision@2"] == pytest.approx(1.0) + assert result["recall@2"] == pytest.approx(5 / 6, abs=1e-3) + + def test_configuration_via_function(self): + """Test configuration via set_numba_enabled function.""" + # Test enabling + ranx.set_numba_enabled(True) + assert ranx.use_numba() is True + + # Test disabling + ranx.set_numba_enabled(False) + assert ranx.use_numba() is False + + def test_same_results_numba_enabled_disabled(self): + """Test that results are the same with Numba enabled vs disabled.""" + qrels_dict = {"q1": {"d1": 2, "d2": 1, "d3": 0}} + run_dict = {"q1": {"d1": 0.9, "d2": 0.8, "d3": 0.7, "d4": 0.6}} + + qrels = ranx.Qrels.from_dict(qrels_dict) + run = ranx.Run.from_dict(run_dict) + + # Test with Numba enabled + ranx.set_numba_enabled(True) + result_numba = ranx.evaluate(qrels, run, ["precision@3", "recall@3", "hits@3"]) + + # Test with Numba disabled + ranx.set_numba_enabled(False) + result_no_numba = ranx.evaluate( + qrels, run, ["precision@3", "recall@3", "hits@3"] + ) + + # Results should be identical + for metric in ["precision@3", "recall@3", "hits@3"]: + assert result_numba[metric] == pytest.approx(result_no_numba[metric])