max-div 0.0.3__py3-none-any.whl → 0.1.1__py3-none-any.whl

max_div/internal/math/fast_log.py

@@ -0,0 +1,167 @@
+import numba
+import numpy as np
+
+# -------------------------------------------------------------------------
+#  Constants
+# -------------------------------------------------------------------------
+
+# --- float64 ---------------------------------------------
+_D_LOG_2 = 0.6931471805599453  # np.log(2)
+
+_D20 = -2.664030016771488
+_D21 = 4.018729576130520
+_D22 = -1.359007257261774
+
+_D30 = -3.144940630924
+_D31 = 6.058956424048
+_D32 = -4.157032648692
+_D33 = 1.243566840636
+
+_D40 = -3.505614661980
+_D41 = 8.099233785172
+_D42 = -8.397609124753
+_D43 = 5.084088932163
+_D44 = -1.280174030509
+
+_D50 = -3.794153676536
+_D51 = 10.139512633266
+_D52 = -14.080875352582
+_D53 = 12.881420375173
+_D54 = -6.551609372263
+_D55 = 1.405716091134
+
+# --- float32 ---------------------------------------------
+_S_LOG_2 = np.float32(_D_LOG_2)
+
+_S20 = np.float32(_D20)
+_S21 = np.float32(_D21)
+_S22 = np.float32(_D22)
+
+_S30 = np.float32(_D30)
+_S31 = np.float32(_D31)
+_S32 = np.float32(_D32)
+_S33 = np.float32(_D33)
+
+_S40 = np.float32(_D40)
+_S41 = np.float32(_D41)
+_S42 = np.float32(_D42)
+_S43 = np.float32(_D43)
+_S44 = np.float32(_D44)
+
+_S50 = np.float32(_D50)
+_S51 = np.float32(_D51)
+_S52 = np.float32(_D52)
+_S53 = np.float32(_D53)
+_S54 = np.float32(_D54)
+_S55 = np.float32(_D55)
+
+
+# -------------------------------------------------------------------------
+#  Fast approximations for np.log
+# -------------------------------------------------------------------------
+@numba.njit(fastmath=True, inline="always")
+def fast_log_f64_poly(x: np.float64, degree: int) -> np.float64:
+    """
+    Fast log approximation using polynomial after range reduction.
+    Accuracy depends on degree:
+        degree=2: max abs error ~0.004   over entire range.
+        degree=3: max abs error ~0.0005  over entire range.
+        degree=4: max abs error ~0.00007 over entire range.
+        degree=5: max abs error ~0.00001 over entire range.
+    """
+    return _D_LOG_2 * fast_log2_f64_poly(x, degree)
+
+
+@numba.njit(fastmath=True, inline="always")
+def fast_log_f32_poly(x: np.float32, degree: int) -> np.float32:
+    """
+    Fast log approximation using polynomial after range reduction.
+    Accuracy depends on degree:
+        degree=2: max abs error ~0.004   over entire range.
+        degree=3: max abs error ~0.0005  over entire range.
+        degree=4: max abs error ~0.00007 over entire range.
+        degree=5: max abs error ~0.00001 over entire range.
+    """
+    return _S_LOG_2 * fast_log2_f32_poly(x, degree)
+
+
+# -------------------------------------------------------------------------
+#  Fast approximations for np.log2
+# -------------------------------------------------------------------------
+@numba.njit(fastmath=True, inline="always")
+def fast_log2_f64_poly(x: np.float64, degree: int) -> np.float64:
+    """
+    Fast log2 approximation using polynomial after range reduction.
+    Accuracy depends on degree:
+        degree=2: max abs error ~0.006    over entire range.
+        degree=3: max abs error ~0.0007   over entire range.
+        degree=4: max abs error ~0.0001   over entire range.
+        degree=5: max abs error ~0.000015 over entire range.
+    """
+
+    # --- extract mantissa & exponent ---------------------
+    # exponent
+    xi = np.int64(np.float64(x).view(np.int64))
+    exponent = ((xi >> 52) & 0x7FF) - 1022
+    # mantissa
+    xi = (xi & 0x000FFFFFFFFFFFFF) | 0x3FE0000000000000
+    m = np.int64(xi).view(np.float64)
+
+    # --- polynomial approximation ------------------------
+    if degree == 2:
+        # log2_mantissa = _D20 + m * (_D21 + m * _D22)
+        log2_mantissa = _D20 + (m * _D21) + (m * m * _D22)
+    elif degree == 3:
+        log2_mantissa = _D30 + m * (_D31 + m * (_D32 + m * _D33))
+    elif degree == 4:
+        log2_mantissa = _D40 + m * (_D41 + m * (_D42 + m * (_D43 + m * _D44)))
+    else:
+        log2_mantissa = _D50 + m * (_D51 + m * (_D52 + m * (_D53 + m * (_D54 + m * _D55))))
+
+    # Return log2(x) = exponent + log2(m)
+    return exponent + log2_mantissa
+
+
+@numba.njit(fastmath=True, inline="always")
+def fast_log2_f32_poly(x: np.float32, degree: int) -> np.float32:
+    """
+    Fast log2 approximation using polynomial after range reduction.
+    Accuracy depends on degree:
+        degree=2: max abs error ~0.006    over entire range.
+        degree=3: max abs error ~0.0007   over entire range.
+        degree=4: max abs error ~0.0001   over entire range.
+        degree=5: max abs error ~0.000015 over entire range.
+    """
+
+    # --- extract mantissa & exponent ---------------------
+    # exponent
+    xi = np.int32(np.float32(x).view(np.int32))
+    exponent = ((xi >> 23) & 0xFF) - 126
+    # mantissa
+    xi = (xi & 0x007FFFFF) | 0x3F000000
+    m = np.int32(xi).view(np.float32)
+
+    # --- polynomial approximation ------------------------
+    if degree == 2:
+        # log2_mantissa = _S20 + m * (_S21 + m * _S22)
+        log2_mantissa = _S20 + (m * _S21) + (m * m * _S22)
+    elif degree == 3:
+        log2_mantissa = _S30 + m * (_S31 + m * (_S32 + m * _S33))
+    elif degree == 4:
+        log2_mantissa = _S40 + m * (_S41 + m * (_S42 + m * (_S43 + m * _S44)))
+    else:
+        log2_mantissa = _S50 + m * (_S51 + m * (_S52 + m * (_S53 + m * (_S54 + m * _S55))))
+
+    # Return log2(x) = exponent + log2(mantissa)
+    return exponent + log2_mantissa
+
+
+# =================================================================================================
+#  Public API
+# =================================================================================================
+__ALL__ = [
+    "fast_log2_f32_poly",
+    "fast_log2_f32_poly",
+    "fast_log_f64_poly",
+    "fast_log_f64_poly",
+]

max_div/internal/math/random.py

@@ -0,0 +1,166 @@
+"""
+Custom simple random number generation module for integration in numba oriented code.  Faster than numpy.random
+even when used inside numba.njit.
+
+This is based on the following:
+  - https://www.pcg-random.org/posts/bounded-rands.html
+  - xoroshiro128+ algorithm by David Blackman and Sebastiano Vigna (http://xoroshiro.di.unimi.it/)
+  - splitmix64 for seed initialization by Sebastiano Vigna (http://xorshift.di.unimi.it/splitmix64.c)
+"""
+
+import numba
+import numpy as np
+from numpy import float32, float64, int32, int64, uint32, uint64
+
+# =================================================================================================
+#  Constants
+# =================================================================================================
+
+# Constant for converting uint64 to float64 in [0.0, 1.0): 1.0 / 2**53
+_TO_FLOAT64 = float64(1.0 / 9007199254740992.0)
+
+
+# Constant for converting uint64 to float32 in [0.0, 1.0): 1.0 / 2**24
+_TO_FLOAT32 = float32(1.0 / 16777216.0)
+
+
+# =================================================================================================
+#  Core
+# =================================================================================================
+@numba.njit(fastmath=True, inline="always")
+def rotl(x: uint64, k: uint64) -> uint64:
+    """Rotate left operation"""
+    return (x << k) | (x >> (uint64(64) - k))
+
+
+@numba.njit(fastmath=True, inline="always")
+def _xoroshiro128plus_next(rng_state: np.ndarray[uint64]) -> uint64:
+    """Generate next random uint64 and update state in-place"""
+    s0 = rng_state[0]
+    s1 = rng_state[1]
+    result = s0 + s1
+
+    s1 ^= s0
+    rng_state[0] = rotl(s0, uint64(24)) ^ s1 ^ (s1 << uint64(16))
+    rng_state[1] = rotl(s1, uint64(37))
+
+    return result
+
+
+@numba.njit(fastmath=True)
+def _splitmix64_next(init_state: np.ndarray[uint64]) -> uint64:
+    """Used to initialize xoroshiro128+ state from single seed; state is a 1-element array, modified in-place."""
+    z = init_state[0] + uint64(0x9E3779B97F4A7C15)
+    init_state[0] = z
+    z = (z ^ (z >> uint64(30))) * uint64(0xBF58476D1CE4E5B9)
+    z = (z ^ (z >> uint64(27))) * uint64(0x94D049BB133111EB)
+    return z ^ (z >> uint64(31))
+
+
+# =================================================================================================
+#  Interface
+# =================================================================================================
+@numba.njit(fastmath=True, inline="always")
+def set_seed(seed: np.int64) -> np.ndarray[uint64]:
+    """Initialize xoroshiro128+ state from single seed; using splitmix64 algorithm."""
+    init_state = np.array([seed], dtype=uint64)
+
+    state = np.empty(2, dtype=uint64)
+    state[0] = _splitmix64_next(init_state)
+    state[1] = _splitmix64_next(init_state)
+
+    return state
+
+
+@numba.njit("float64(uint64[:])", fastmath=True, inline="always")
+def rand_float64(rng_state: np.ndarray[uint64]) -> float64:
+    """Generate a random float64 in [0.0, 1.0) using the provided rng_state."""
+    rnd_uint64 = _xoroshiro128plus_next(rng_state)
+    return float64((rnd_uint64 >> uint64(11)) * _TO_FLOAT64)  # 2**-53
+
+
+@numba.njit("float32(uint64[:])", fastmath=True, inline="always")
+def rand_float32(rng_state: np.ndarray[uint64]) -> float32:
+    """Generate a random float32 in [0.0, 1.0) using the provided rng_state."""
+    rnd_uint64 = _xoroshiro128plus_next(rng_state)
+    return float32((rnd_uint64 >> uint64(40)) * _TO_FLOAT32)  # 2**-24
+
+
+@numba.njit("int64(uint64[:], int64, int64)", fastmath=True, inline="always")
+def rand_int64(rng_state: np.ndarray[uint64], low: np.int64, high: np.int64) -> np.int64:
+    """
+    Generate a random int64 in [low, high) using the provided rng_state.
+    There might be a small bias for large (high-low) if the range is not a power of two.
+    """
+    if low == 0:
+        rnd_uint64 = _xoroshiro128plus_next(rng_state)
+        return int64(rnd_uint64 % uint64(high))
+    else:
+        range_size = high - low
+        rnd_uint64 = _xoroshiro128plus_next(rng_state)
+        return low + int64(rnd_uint64 % uint64(range_size))
+
+
+@numba.njit("int32(uint64[:], int32, int32)", fastmath=True, inline="always")
+def rand_int32(rng_state: np.ndarray[uint64], low: np.int32, high: np.int32) -> np.int32:
+    """
+    Generate a random int32 in [low, high) using the provided rng_state.
+    There might be a small bias for large (high-low) if the range is not a power of two.
+    """
+    if low == 0:
+        rnd_uint64 = _xoroshiro128plus_next(rng_state)
+        return int32(rnd_uint64 % uint64(high))
+    else:
+        range_size = high - low
+        rnd_uint64 = _xoroshiro128plus_next(rng_state)
+        return low + int32(rnd_uint64 % uint64(range_size))
+
+
+@numba.njit("int32[:](uint64[:], int32, int32, int32)", fastmath=True, inline="always")
+def rand_int32_array(
+    rng_state: np.ndarray[uint64], low: np.int32, high: np.int32, size: np.int32
+) -> np.ndarray[np.int32]:
+    """
+    Generate an array of random int32 values in [low, high) using the provided rng_state.
+    Optimized to generate 2 values per RNG call by using upper and lower 32 bits.
+    There might be a small bias for large (high-low) if the range is not a power of two.
+    """
+    result = np.empty(size, dtype=np.int32)
+    if low == 0:
+        range_size = uint64(high)
+        i = 0
+        while i < size:
+            rnd_uint64 = _xoroshiro128plus_next(rng_state)
+            # Use lower 32 bits for first value
+            result[i] = int32((rnd_uint64 & uint64(0xFFFFFFFF)) % range_size)
+            i += 1
+            # Use upper 32 bits for second value if needed
+            if i < size:
+                result[i] = int32((rnd_uint64 >> uint64(32)) % range_size)
+                i += 1
+    else:
+        range_size = uint64(high - low)
+        i = 0
+        while i < size:
+            rnd_uint64 = _xoroshiro128plus_next(rng_state)
+            # Use lower 32 bits for first value
+            result[i] = low + int32((rnd_uint64 & uint64(0xFFFFFFFF)) % range_size)
+            i += 1
+            # Use upper 32 bits for second value if needed
+            if i < size:
+                result[i] = low + int32((rnd_uint64 >> uint64(32)) % range_size)
+                i += 1
+    return result
+
+
+# =================================================================================================
+#  Public API
+# =================================================================================================
+__ALL__ = [
+    "rand_float32",
+    "rand_float64",
+    "rand_int32",
+    "rand_int32_array",
+    "rand_int64",
+    "set_seed",
+]

max_div/internal/math/select_k_minmax.py

@@ -0,0 +1,250 @@
+import numba
+import numpy as np
+
+
+# =================================================================================================
+#  select_k_min
+# =================================================================================================
+@numba.njit("int32[:](float32[:], int32)", fastmath=True, inline="always")
+def select_k_min(arr: np.ndarray[np.float32], k: np.int32) -> np.ndarray[np.int32]:
+    """
+    Find indices of k smallest elements in a float32 array using Numba.
+
+    This implementation uses a max-heap approach with O(n log k) complexity,
+    which is efficient when k << n. The heap maintains the k smallest elements
+    seen so far, with the largest of these at the root.
+
+    Parameters:
+    -----------
+    arr : np.ndarray[np.float32]
+        Input array with n elements (typically 1000-10000)
+    k : int
+        Number of smallest elements to find
+
+    Returns:
+    --------
+    indices : np.ndarray[np.int32]
+        Array of k indices pointing to the smallest elements.
+        Indices are returned in arbitrary order (not sorted by value).
+
+    Performance:
+    ------------
+    - 2-8x faster than np.argpartition for small to moderate k (k ~ 10-100)
+    - Best when k << n (e.g., k=100, n=10000)
+    - Uses fastmath=True for additional SIMD optimizations
+    """
+    n = len(arr)
+    heap_idx = np.empty(k, dtype=np.int32)  # indices (into arr) of elements in the heap
+    heap_values = np.empty(k, dtype=np.float32)  # values of elements in the heap; largest at heap_values[0]
+
+    # Build initial heap with first k elements
+    for i in range(k):
+        heap_idx[i] = i
+        heap_values[i] = arr[i]
+
+    # Heapify: Convert initial k elements into a max-heap
+    # -----------------------------------------------------------------------------------
+    #
+    # assuming we want to represent values v0 >= v1 >= v2 >= v3 >= v4 >= v5 >= v6 into a heap:
+    #
+    #            v0
+    #          /    \
+    #        v1      v2
+    #       / \     / \
+    #     v3  v4  v5  v6
+    #
+    # Invariant relations:
+    #   - parents >= leaves   (i.e. v0 >= v1,v2; v1 >= v3,v4; v2 >= v5,v6)
+    #   - leaves of same parent are not necessarily sorted (!)  (i.e. the tree could swap branches v1 & v2)
+    #   - if a parent is at index i, then its children are at indices 2*i+1 and 2*i+2
+    # -----------------------------------------------------------------------------------
+    # Start from last non-leaf node and sift down
+    for i in range(k // 2 - 1, -1, -1):
+        i_parent = i
+        value = heap_values[i_parent]
+        idx = heap_idx[i_parent]
+
+        # Sift down: move element down until heap property is restored
+        while True:
+            i_child_left = 2 * i_parent + 1
+            i_child_right = i_child_left + 1
+            i_child_largest = -1
+
+            # Find the largest child
+            if i_child_left < k:
+                if i_child_right < k:
+                    i_child_largest = (
+                        i_child_left if heap_values[i_child_left] > heap_values[i_child_right] else i_child_right
+                    )
+                else:
+                    i_child_largest = i_child_left
+
+            # If no children or value is larger than the largest child, we're done
+            if i_child_largest == -1 or value >= heap_values[i_child_largest]:
+                heap_values[i_parent] = value
+                heap_idx[i_parent] = idx
+                break
+
+            # Otherwise, move the larger child up and continue
+            heap_values[i_parent] = heap_values[i_child_largest]
+            heap_idx[i_parent] = heap_idx[i_child_largest]
+            i_parent = i_child_largest
+
+    # Process remaining elements
+    # For each element, if it's smaller than heap maximum, replace and sift down
+    for i in range(k, n):
+        value = arr[i]
+        if value < heap_values[0]:  # heap_values[0] is the maximum of k smallest
+            i_parent = 0
+
+            # Sift down from root
+            while True:
+                i_child_left = 2 * i_parent + 1
+                i_child_right = i_child_left + 1
+                i_child_largest = -1
+
+                # Find the largest child
+                if i_child_left < k:
+                    if i_child_right < k:
+                        i_child_largest = (
+                            i_child_left if heap_values[i_child_left] > heap_values[i_child_right] else i_child_right
+                        )
+                    else:
+                        i_child_largest = i_child_left
+
+                # If no children or val is larger than largest child, we're done
+                if i_child_largest == -1 or value >= heap_values[i_child_largest]:
+                    heap_values[i_parent] = value
+                    heap_idx[i_parent] = i
+                    break
+
+                # Otherwise, move the larger child up and continue
+                heap_values[i_parent] = heap_values[i_child_largest]
+                heap_idx[i_parent] = heap_idx[i_child_largest]
+                i_parent = i_child_largest
+
+    return heap_idx
+
+
+# =================================================================================================
+#  select_k_max
+# =================================================================================================
+@numba.njit("int32[:](float32[:], int32)", fastmath=True, inline="always")
+def select_k_max(arr: np.ndarray[np.float32], k: np.int32) -> np.ndarray[np.int32]:
+    """
+    Find indices of k largest elements in a float32 array using Numba.
+
+    This implementation uses a min-heap approach with O(n log k) complexity,
+    which is efficient when k << n. The heap maintains the k largest elements
+    seen so far, with the smallest of these at the root.
+
+    Parameters:
+    -----------
+    arr : np.ndarray[np.float32]
+        Input array with n elements (typically 1000-10000)
+    k : int
+        Number of largest elements to find
+
+    Returns:
+    --------
+    indices : np.ndarray[np.int32]
+        Array of k indices pointing to the largest elements.
+        Indices are returned in arbitrary order (not sorted by value).
+
+    Performance:
+    ------------
+    - 2-8x faster than np.argpartition for small to moderate k (k ~ 10-100)
+    - Best when k << n (e.g., k=100, n=10000)
+    - Uses fastmath=True for additional SIMD optimizations
+    """
+    n = len(arr)
+    heap_idx = np.empty(k, dtype=np.int32)  # indices (into arr) of elements in the heap
+    heap_values = np.empty(k, dtype=np.float32)  # values of elements in the heap; smallest at heap_values[0]
+
+    # Build initial heap with first k elements
+    for i in range(k):
+        heap_idx[i] = i
+        heap_values[i] = arr[i]
+
+    # Heapify: Convert initial k elements into a min-heap
+    # -----------------------------------------------------------------------------------
+    #
+    # assuming we want to represent values v0 <= v1 <= v2 <= v3 <= v4 <= v5 <= v6 into a heap:
+    #
+    #            v0
+    #          /    \
+    #        v1      v2
+    #       / \     / \
+    #     v3  v4  v5  v6
+    #
+    # Invariant relations:
+    #   - parents <= leaves   (i.e. v0 <= v1,v2; v1 <= v3,v4; v2 <= v5,v6)
+    #   - leaves of same parent are not necessarily sorted (!)  (i.e. the tree could swap branches v1 & v2)
+    #   - if a parent is at index i, then its children are at indices 2*i+1 and 2*i+2
+    # -----------------------------------------------------------------------------------
+    # Start from last non-leaf node and sift down
+    for i in range(k // 2 - 1, -1, -1):
+        i_parent = i
+        value = heap_values[i_parent]
+        idx = heap_idx[i_parent]
+
+        # Sift down: move element down until heap property is restored
+        while True:
+            i_child_left = 2 * i_parent + 1
+            i_child_right = i_child_left + 1
+            i_child_smallest = -1
+
+            # Find the smallest child
+            if i_child_left < k:
+                if i_child_right < k:
+                    i_child_smallest = (
+                        i_child_left if heap_values[i_child_left] < heap_values[i_child_right] else i_child_right
+                    )
+                else:
+                    i_child_smallest = i_child_left
+
+            # If no children or value is smaller than the smallest child, we're done
+            if i_child_smallest == -1 or value <= heap_values[i_child_smallest]:
+                heap_values[i_parent] = value
+                heap_idx[i_parent] = idx
+                break
+
+            # Otherwise, move the smaller child up and continue
+            heap_values[i_parent] = heap_values[i_child_smallest]
+            heap_idx[i_parent] = heap_idx[i_child_smallest]
+            i_parent = i_child_smallest
+
+    # Process remaining elements
+    # For each element, if it's larger than heap minimum, replace and sift down
+    for i in range(k, n):
+        value = arr[i]
+        if value > heap_values[0]:  # heap_values[0] is the minimum of k largest
+            i_parent = 0
+
+            # Sift down from root
+            while True:
+                i_child_left = 2 * i_parent + 1
+                i_child_right = i_child_left + 1
+                i_child_smallest = -1
+
+                # Find the smallest child
+                if i_child_left < k:
+                    if i_child_right < k:
+                        i_child_smallest = (
+                            i_child_left if heap_values[i_child_left] < heap_values[i_child_right] else i_child_right
+                        )
+                    else:
+                        i_child_smallest = i_child_left
+
+                # If no children or val is smaller than smallest child, we're done
+                if i_child_smallest == -1 or value <= heap_values[i_child_smallest]:
+                    heap_values[i_parent] = value
+                    heap_idx[i_parent] = i
+                    break
+
+                # Otherwise, move the smaller child up and continue
+                heap_values[i_parent] = heap_values[i_child_smallest]
+                heap_idx[i_parent] = heap_idx[i_child_smallest]
+                i_parent = i_child_smallest
+
+    return heap_idx

max_div/sampling/__init__.py

@@ -1 +1 @@
-from .discrete import sample_int
+from .uncon import randint, randint_numba, randint_numpy