numba-utils 0.1.0__py3-none-any.whl

numba_utils/__init__.py

@@ -0,0 +1,137 @@
+"""numba-utils: high-performance building blocks for Numba."""
+
+from numba_utils import diagnostics, testing
+from numba_utils._config import config, configure
+from numba_utils.algorithms import (
+    argmax2,
+    counting_sort,
+    fast_argpartition,
+    insertion_sort,
+    nth_element,
+    partial_sort,
+    quickselect,
+    radix_sort,
+    topk,
+)
+from numba_utils.arrays import (
+    bincount,
+    binary_search,
+    cumulative_sum,
+    fast_clip,
+    histogram,
+    lower_bound,
+    normalize,
+    rolling_mean,
+    rolling_sum,
+    unique_sorted,
+    upper_bound,
+)
+from numba_utils.collections import (
+    BitSet,
+    FixedQueue,
+    ObjectPool,
+    PriorityQueue,
+    RingBuffer,
+    SparseSet,
+    Stack,
+    counter,
+    typed_defaultdict,
+)
+from numba_utils.decorators import (
+    boundscheck,
+    cached_njit,
+    njit_fast,
+    njit_parallel,
+)
+from numba_utils.parallel import (
+    parallel_histogram,
+    parallel_prefix_sum,
+    parallel_reduce,
+    parallel_sum,
+    parallel_topk,
+)
+from numba_utils.profiling import (
+    BenchmarkResult,
+    ComparisonResult,
+    TimingStats,
+    benchmark,
+    compare,
+    compile_stats,
+    compile_time,
+    warmup,
+)
+from numba_utils.random import (
+    alias_draw,
+    alias_sample,
+    alias_setup,
+    choice,
+    permutation,
+    reservoir_sampling,
+    seed,
+    shuffle,
+    weighted_sampling,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "BenchmarkResult",
+    "BitSet",
+    "ComparisonResult",
+    "FixedQueue",
+    "ObjectPool",
+    "PriorityQueue",
+    "RingBuffer",
+    "SparseSet",
+    "Stack",
+    "TimingStats",
+    "alias_draw",
+    "alias_sample",
+    "alias_setup",
+    "argmax2",
+    "benchmark",
+    "bincount",
+    "binary_search",
+    "boundscheck",
+    "cached_njit",
+    "choice",
+    "compare",
+    "compile_stats",
+    "compile_time",
+    "config",
+    "configure",
+    "counter",
+    "counting_sort",
+    "cumulative_sum",
+    "diagnostics",
+    "fast_argpartition",
+    "fast_clip",
+    "histogram",
+    "insertion_sort",
+    "lower_bound",
+    "njit_fast",
+    "njit_parallel",
+    "normalize",
+    "nth_element",
+    "parallel_histogram",
+    "parallel_prefix_sum",
+    "parallel_reduce",
+    "parallel_sum",
+    "parallel_topk",
+    "partial_sort",
+    "permutation",
+    "quickselect",
+    "radix_sort",
+    "reservoir_sampling",
+    "rolling_mean",
+    "rolling_sum",
+    "seed",
+    "shuffle",
+    "testing",
+    "topk",
+    "typed_defaultdict",
+    "unique_sorted",
+    "upper_bound",
+    "warmup",
+    "weighted_sampling",
+]

numba_utils/_config.py

@@ -0,0 +1,92 @@
+"""Global configuration: code-level overrides with environment fallbacks.
+
+Two levels, resolved per option at decoration time:
+
+1. Code: ``configure(cache=False)`` or ``config.cache = False``.
+2. Environment: ``NUMBA_UTILS_CACHE=0`` (checked only when the code level
+   is unset).
+
+A set option (``True``/``False``) is a GLOBAL override: it wins over
+decorator defaults and per-call keyword arguments alike. That is the
+point — overrides exist for environment-level policy, like disabling the
+on-disk cache on machines where it is unsafe (see docs/numba-cache.md).
+``None`` means "no override": decorator defaults and per-call arguments
+apply unchanged.
+
+Overrides affect only functions decorated AFTER the change — configure
+before defining your jitted functions.
+"""
+
+from __future__ import annotations
+
+import os
+
+_TRUTHY = frozenset({"1", "true", "yes", "on"})
+_FALSY = frozenset({"0", "false", "no", "off"})
+
+OPTION_ENV_VARS = {
+    "cache": "NUMBA_UTILS_CACHE",
+    "fastmath": "NUMBA_UTILS_FASTMATH",
+    "parallel": "NUMBA_UTILS_PARALLEL",
+    "nogil": "NUMBA_UTILS_NOGIL",
+}
+
+
+class GlobalConfig:
+    """Tri-state global overrides for every numba-utils decorator.
+
+    Attributes ``cache``, ``fastmath``, ``parallel``, ``nogil`` are each
+    ``None`` (no override), ``True`` or ``False`` (forced globally).
+    """
+
+    __slots__ = ("cache", "fastmath", "parallel", "nogil")
+
+    def __init__(self) -> None:
+        self.reset()
+
+    def reset(self) -> None:
+        """Clear all code-level overrides. Environment fallbacks remain."""
+        for name in OPTION_ENV_VARS:
+            setattr(self, name, None)
+
+    def resolve(self, name: str) -> bool | None:
+        """Effective override for ``name``: code level, else environment,
+        else ``None``. Unrecognized environment values count as unset."""
+        value = getattr(self, name)
+        if value is not None:
+            return bool(value)
+        raw = os.environ.get(OPTION_ENV_VARS[name], "").strip().lower()
+        if raw in _TRUTHY:
+            return True
+        if raw in _FALSY:
+            return False
+        return None
+
+
+config = GlobalConfig()
+
+
+def configure(**options: bool | None) -> None:
+    """Set global decorator overrides in code.
+
+    ::
+
+        import numba_utils as nu
+        nu.configure(cache=False, fastmath=True)
+
+    Valid options: ``cache``, ``fastmath``, ``parallel``, ``nogil``.
+    Pass ``None`` to clear a single override; ``config.reset()`` clears
+    all. Raises ``ValueError`` on unknown names (fail fast, no typos
+    silently ignored).
+    """
+    for name, value in options.items():
+        if name not in OPTION_ENV_VARS:
+            valid = ", ".join(sorted(OPTION_ENV_VARS))
+            raise ValueError(
+                f"configure: unknown option {name!r}; valid options: {valid}"
+            )
+        if value is not None and not isinstance(value, bool):
+            raise ValueError(
+                f"configure: option {name!r} must be True, False or None"
+            )
+        setattr(config, name, value)

numba_utils/algorithms/__init__.py

@@ -0,0 +1,26 @@
+"""Selection and sorting algorithms specialized for Numba."""
+
+from numba_utils.algorithms._select import (
+    fast_argpartition,
+    nth_element,
+    quickselect,
+)
+from numba_utils.algorithms._sort import (
+    counting_sort,
+    insertion_sort,
+    partial_sort,
+    radix_sort,
+)
+from numba_utils.algorithms._topk import argmax2, topk
+
+__all__ = [
+    "argmax2",
+    "counting_sort",
+    "fast_argpartition",
+    "insertion_sort",
+    "nth_element",
+    "partial_sort",
+    "quickselect",
+    "radix_sort",
+    "topk",
+]

numba_utils/algorithms/_select.py

@@ -0,0 +1,123 @@
+"""Selection: quickselect, nth_element and argpartition.
+
+Mutation contract, stated per function and consistent across the module:
+``nth_element`` partitions IN PLACE (the zero-allocation building block);
+``quickselect`` and ``fast_argpartition`` never touch their input.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from numba_utils.decorators import cached_njit
+
+
+@cached_njit
+def _median3(a, b, c):
+    if a < b:
+        if b < c:
+            return b
+        elif a < c:
+            return c
+        else:
+            return a
+    else:
+        if a < c:
+            return a
+        elif b < c:
+            return c
+        else:
+            return b
+
+
+@cached_njit
+def nth_element(arr, k):
+    """Partition ``arr`` IN PLACE so ``arr[k]`` is its k-th smallest value.
+
+    After the call: ``arr[:k] <= arr[k] <= arr[k+1:]`` (both sides
+    unordered). Returns ``arr[k]``. Iterative Hoare quickselect with
+    median-of-three pivots. This is the C++ ``std::nth_element``.
+
+    Complexity: average O(n), worst O(n²). Memory: O(1).
+    """
+    n = arr.shape[0]
+    if n == 0:
+        raise ValueError("nth_element: empty array")
+    if k < 0 or k >= n:
+        raise ValueError("nth_element: k out of range")
+    lo = 0
+    hi = n - 1
+    while lo < hi:
+        mid = (lo + hi) >> 1
+        pivot = _median3(arr[lo], arr[mid], arr[hi])
+        i = lo
+        j = hi
+        while i <= j:
+            while arr[i] < pivot:
+                i += 1
+            while arr[j] > pivot:
+                j -= 1
+            if i <= j:
+                arr[i], arr[j] = arr[j], arr[i]
+                i += 1
+                j -= 1
+        if k <= j:
+            hi = j
+        elif k >= i:
+            lo = i
+        else:
+            break
+    return arr[k]
+
+
+@cached_njit
+def quickselect(arr, k):
+    """k-th smallest value of ``arr`` (0-indexed). Input is NOT modified.
+
+    Convenience wrapper: copies, then runs :func:`nth_element`. Call
+    ``nth_element`` directly in hot loops to skip the copy.
+
+    Complexity: average O(n), worst O(n²). Memory: O(n) for the copy.
+    """
+    tmp = arr.copy()
+    return nth_element(tmp, k)
+
+
+@cached_njit
+def fast_argpartition(arr, k):
+    """Indices of the k smallest elements of ``arr``, in no particular order.
+
+    Like ``np.argpartition(arr, k)[:k]`` but returns exactly k indices and
+    skips NumPy's full output. Input is NOT modified; the quickselect runs
+    on an index array compared through ``arr``.
+
+    Complexity: average O(n), worst O(n²). Memory: O(n) for the indices.
+    """
+    n = arr.shape[0]
+    if k < 1 or k > n:
+        raise ValueError("fast_argpartition: k must be in [1, len(arr)]")
+    idx = np.arange(n)
+    target = k - 1
+    lo = 0
+    hi = n - 1
+    while lo < hi:
+        mid = (lo + hi) >> 1
+        pivot = _median3(arr[idx[lo]], arr[idx[mid]], arr[idx[hi]])
+        i = lo
+        j = hi
+        while i <= j:
+            while arr[idx[i]] < pivot:
+                i += 1
+            while arr[idx[j]] > pivot:
+                j -= 1
+            if i <= j:
+                idx[i], idx[j] = idx[j], idx[i]
+                i += 1
+                j -= 1
+        if target <= j:
+            hi = j
+        elif target >= i:
+            lo = i
+        else:
+            break
+    return idx[:k].copy()

numba_utils/algorithms/_sort.py

@@ -0,0 +1,166 @@
+"""Sorting: insertion sort, partial sort, counting sort, radix sort.
+
+Mutation contract: ``insertion_sort`` and ``partial_sort`` work IN PLACE
+(zero allocation — that's their reason to exist) and return the array for
+chaining. ``counting_sort`` and ``radix_sort`` allocate scratch space
+inherently, so they return a NEW array and leave the input untouched.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from numba_utils.algorithms._select import nth_element
+from numba_utils.decorators import cached_njit
+
+_COUNTING_SORT_MAX_RANGE = 2**27
+
+
+@cached_njit
+def insertion_sort(arr):
+    """Sort ``arr`` ascending IN PLACE; returns it.
+
+    O(n²), but with tiny constants and O(n) on nearly-sorted input —
+    the right tool below ~64 elements or as a small-partition finisher.
+
+    Complexity: O(n²) worst, O(n) nearly-sorted. Memory: O(1).
+    """
+    for i in range(1, arr.shape[0]):
+        x = arr[i]
+        j = i - 1
+        while j >= 0 and arr[j] > x:
+            arr[j + 1] = arr[j]
+            j -= 1
+        arr[j + 1] = x
+    return arr
+
+
+@cached_njit
+def partial_sort(arr, k):
+    """Rearrange ``arr`` IN PLACE so ``arr[:k]`` holds the k smallest, sorted.
+
+    ``arr[k:]`` contains the rest in arbitrary order. Returns ``arr``.
+    The C++ ``std::partial_sort``: select with :func:`nth_element`, then
+    sort only the front.
+
+    Complexity: average O(n + k log k). Memory: O(1).
+    """
+    n = arr.shape[0]
+    if k < 1 or k > n:
+        raise ValueError("partial_sort: k must be in [1, len(arr)]")
+    if k < n:
+        nth_element(arr, k - 1)
+    arr[:k].sort()
+    return arr
+
+
+@cached_njit
+def counting_sort(arr):
+    """Return a NEW sorted array of integers via counting sort.
+
+    Non-comparison sort: O(n + range) where range = max - min + 1. Wins
+    when the value range is small relative to n (codes, buckets, bytes).
+    Raises ``ValueError`` if range exceeds 2**27 (a 1 GiB counts array) —
+    use :func:`radix_sort` there. Integer dtypes only.
+
+    Complexity: O(n + range). Memory: O(n + range).
+    """
+    n = arr.shape[0]
+    out = np.empty_like(arr)
+    if n == 0:
+        return out
+    mn = arr[0]
+    mx = arr[0]
+    for i in range(1, n):
+        v = arr[i]
+        if v < mn:
+            mn = v
+        elif v > mx:
+            mx = v
+    value_range = np.int64(mx) - np.int64(mn) + 1
+    if value_range > _COUNTING_SORT_MAX_RANGE:
+        raise ValueError(
+            "counting_sort: value range too large (> 2**27), use radix_sort"
+        )
+    counts = np.zeros(value_range, np.int64)
+    for i in range(n):
+        counts[np.int64(arr[i]) - np.int64(mn)] += 1
+    j = 0
+    for v in range(value_range):
+        value = mn + v
+        for _ in range(counts[v]):
+            out[j] = value
+            j += 1
+    return out
+
+
+@cached_njit
+def _radix_key(x, mn):
+    # Modular difference (x - mn) mod 2**64 equals the true nonnegative
+    # distance for every integer dtype, so one biased uint64 key handles
+    # signed and unsigned inputs alike.
+    return np.uint64(np.int64(x) - np.int64(mn))
+
+
+@cached_njit
+def radix_sort(arr):
+    """Return a NEW sorted array of integers via LSD radix sort.
+
+    8-bit digits over min-biased uint64 keys; handles all signed and
+    unsigned integer dtypes, negatives included. All digit histograms are
+    built in ONE pass, then each byte gets a scatter pass — and bytes
+    that are constant across the array (including the high bytes the
+    value range doesn't reach) are skipped entirely. Integer dtypes only.
+
+    When it wins: values spanning <= 3-4 bytes (each byte is a full
+    scatter pass). For full-range 64-bit keys, NumPy's SIMD introsort
+    is faster — see BENCHMARKS.md for both cases.
+
+    Complexity: O(n · (1 + bytes_used)). Memory: O(n) scratch + counters.
+    """
+    n = arr.shape[0]
+    a = arr.copy()
+    if n <= 1:
+        return a
+    b = np.empty_like(arr)
+    mn = arr[0]
+    mx = arr[0]
+    for i in range(1, n):
+        v = arr[i]
+        if v < mn:
+            mn = v
+        elif v > mx:
+            mx = v
+    max_key = _radix_key(mx, mn)
+    eight = np.uint64(8)
+    mask = np.uint64(0xFF)
+    passes = 1
+    while passes < 8 and (max_key >> np.uint64(8 * passes)) != np.uint64(0):
+        passes += 1
+    counts = np.zeros((passes, 256), np.int64)
+    for i in range(n):
+        key = _radix_key(a[i], mn)
+        for p in range(passes):
+            counts[p, key & mask] += 1
+            key = key >> eight
+    for p in range(passes):
+        shift = np.uint64(8 * p)
+        cp = counts[p]
+        constant_byte = False
+        for d in range(256):
+            if cp[d] == n:
+                constant_byte = True
+                break
+        if constant_byte:
+            continue
+        total = 0
+        for d in range(256):
+            c = cp[d]
+            cp[d] = total
+            total += c
+        for i in range(n):
+            digit = (_radix_key(a[i], mn) >> shift) & mask
+            b[cp[digit]] = a[i]
+            cp[digit] += 1
+        a, b = b, a
+    return a

numba_utils/algorithms/_topk.py

@@ -0,0 +1,87 @@
+"""Top-k and extremum queries."""
+
+from __future__ import annotations
+
+from numba_utils.algorithms._select import nth_element
+from numba_utils.decorators import cached_njit
+
+# Below n/16, a size-k min-heap beats quickselect: one read-only pass,
+# no O(n) copy, and most elements fail the `x > heap[0]` test immediately.
+_HEAP_PATH_FACTOR = 16
+
+
+@cached_njit
+def _sift_down(heap, start, size):
+    root = start
+    while True:
+        child = 2 * root + 1
+        if child >= size:
+            return
+        if child + 1 < size and heap[child + 1] < heap[child]:
+            child += 1
+        if heap[child] < heap[root]:
+            heap[root], heap[child] = heap[child], heap[root]
+            root = child
+        else:
+            return
+
+
+@cached_njit
+def _topk_heap(arr, k):
+    n = arr.shape[0]
+    heap = arr[:k].copy()
+    for start in range(k // 2 - 1, -1, -1):
+        _sift_down(heap, start, k)
+    for i in range(k, n):
+        x = arr[i]
+        if x > heap[0]:
+            heap[0] = x
+            _sift_down(heap, 0, k)
+    heap.sort()
+    return heap[::-1].copy()
+
+
+@cached_njit
+def topk(arr, k):
+    """The k LARGEST values of ``arr``, sorted descending. Input untouched.
+
+    Small k (``k * 16 <= n``): single read-only pass with a size-k
+    min-heap, O(n) time and O(k) extra memory. Large k: quickselect via
+    :func:`nth_element` on a copy, then sort only the k winners. Either
+    way the full sort never happens. For the k smallest, use
+    :func:`fast_argpartition`.
+
+    Complexity: average O(n + k log k). Memory: O(k) small k, O(n) large k.
+    """
+    n = arr.shape[0]
+    if k < 1 or k > n:
+        raise ValueError("topk: k must be in [1, len(arr)]")
+    if k * _HEAP_PATH_FACTOR <= n:
+        return _topk_heap(arr, k)
+    tmp = arr.copy()
+    if k < n:
+        nth_element(tmp, n - k)
+    result = tmp[n - k :]
+    result.sort()
+    return result[::-1].copy()
+
+
+@cached_njit
+def argmax2(arr):
+    """Index AND value of the maximum, in one pass: ``(idx, value)``.
+
+    Saves the second scan of the ``arr[np.argmax(arr)]`` idiom. First
+    occurrence wins on ties. Raises ``ValueError`` on empty input.
+
+    Complexity: O(n). Memory: O(1).
+    """
+    n = arr.shape[0]
+    if n == 0:
+        raise ValueError("argmax2: empty array")
+    best = arr[0]
+    best_idx = 0
+    for i in range(1, n):
+        if arr[i] > best:
+            best = arr[i]
+            best_idx = i
+    return best_idx, best

numba_utils/arrays/__init__.py

@@ -0,0 +1,21 @@
+"""Array utilities: searching, transforms, rolling windows, histograms."""
+
+from numba_utils.arrays._hist import bincount, histogram
+from numba_utils.arrays._rolling import rolling_mean, rolling_sum
+from numba_utils.arrays._search import binary_search, lower_bound, upper_bound
+from numba_utils.arrays._transform import cumulative_sum, fast_clip, normalize
+from numba_utils.arrays._unique import unique_sorted
+
+__all__ = [
+    "bincount",
+    "binary_search",
+    "cumulative_sum",
+    "fast_clip",
+    "histogram",
+    "lower_bound",
+    "normalize",
+    "rolling_mean",
+    "rolling_sum",
+    "unique_sorted",
+    "upper_bound",
+]

numba_utils/arrays/_hist.py

@@ -0,0 +1,56 @@
+"""Integer-optimized counting: bincount and fixed-range histogram."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from numba_utils.decorators import cached_njit
+
+
+@cached_njit
+def bincount(arr, minlength=0):
+    """Count occurrences of each value in a 1-D array of nonnegative ints.
+
+    ``result[v]`` is how many times ``v`` appears. Output length is
+    ``max(arr.max() + 1, minlength)``. Raises ``ValueError`` on negatives.
+
+    Complexity: O(n + max). Memory: O(max).
+    """
+    mx = minlength - 1
+    for i in range(arr.shape[0]):
+        v = arr[i]
+        if v < 0:
+            raise ValueError("bincount: negative values are not allowed")
+        if v > mx:
+            mx = v
+    counts = np.zeros(mx + 1, np.int64)
+    for i in range(arr.shape[0]):
+        counts[arr[i]] += 1
+    return counts
+
+
+@cached_njit
+def histogram(arr, bins, lo, hi):
+    """Histogram of 1-D ``arr`` over ``bins`` equal-width bins in ``[lo, hi]``.
+
+    Single pass, no edge array: bin index is computed by scaling, which is
+    why this beats ``np.histogram``. Values outside ``[lo, hi]`` are
+    ignored; ``hi`` itself lands in the last bin (NumPy convention).
+
+    Complexity: O(n + bins). Memory: O(bins).
+    """
+    if bins < 1:
+        raise ValueError("histogram: bins must be >= 1")
+    if not lo < hi:
+        raise ValueError("histogram: lo must be < hi")
+    counts = np.zeros(bins, np.int64)
+    scale = bins / (hi - lo)
+    for i in range(arr.shape[0]):
+        x = arr[i]
+        if x < lo or x > hi:
+            continue
+        idx = int((x - lo) * scale)
+        if idx >= bins:
+            idx = bins - 1
+        counts[idx] += 1
+    return counts