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

max_div/benchmark/randint_constrained.py

@@ -0,0 +1,355 @@
+import math
+from dataclasses import dataclass
+
+import numpy as np
+from tqdm import tqdm
+
+from max_div.constraints._numba import _build_array_repr
+from max_div.internal.benchmarking import BenchmarkResult, benchmark
+from max_div.internal.formatting import md_multiline
+from max_div.sampling import randint_numba
+from max_div.sampling.con import Constraint, randint_constrained_numba
+
+from ._formatting import (
+    BoldLabels,
+    CellContent,
+    FastestBenchmark,
+    HighestPercentage,
+    Percentage,
+    extend_table_with_aggregate_row,
+    format_as_markdown,
+    format_for_console,
+)
+
+
+# =================================================================================================
+#  Helper classes
+# =================================================================================================
+@dataclass
+class _Scenario:
+    letter: str
+    description: str
+    n_k_cons_tuples: list[tuple[int, int, int]]  # list of (n, k, n_cons)-tuples
+    use_p: bool
+
+
+# =================================================================================================
+#  Main benchmark function
+# =================================================================================================
+def benchmark_randint_constrained(speed: float = 0.0, markdown: bool = False) -> None:
+    """
+    Benchmarks the `randint_constrained` function from `max_div.sampling.con`.
+
+    Different scenarios are tested across different values of `k`, `n` & `n_cons` (# of constraints):
+
+     * **SCENARIO A1**
+        * all combinations with `k` < `n` with
+            * `n` in [10, 100, 1000]
+            * `k` in [2, 4, 8, 16, 32, ..., 256]
+        * constraints:
+            * 10 non-overlapping constraints, each spanning exactly 1/10th of the `n`-range
+            * min_count = floor(k/11)
+            * max_count = ceil(k/9)
+        * no `p` provided (uniform sampling)
+
+     * **SCENARIO A2**
+        * same as Scenario A1, but with `p` provided (with p[i] ~ 1+i)
+
+     * **SCENARIO B1**
+        * `n` =  1000
+        * `k` =   100
+        * `n_cons` in [2, 4, 8, 16, ..., 512]
+            * each constraint spans a random 1% of the `n` range (=10 values)
+            * min_count = floor(10 / n_cons)
+            * max_count = ceil(100 / n_cons)
+        * no `p` provided (uniform sampling)
+
+     * **SCENARIO B2**
+        * same as Scenario B1, but with `p` provided (with p[i] ~ 1+i)
+
+    :param speed: value in [0.0, 1.0] (default=0.0); 0.0=accurate but slow; 1.0=fast but less accurate
+    :param markdown: If `True`, outputs the results as a Markdown table.
+    """
+
+    # --- build scenarios ---------------------------------
+    scenarios = []
+    for use_p in [False, True]:
+        if not use_p:
+            letter = "A1"
+            description = "Varying n & k with 10 non-overlapping constraints spanning equal portions of the n range (uniform sampling)."
+        else:
+            letter = "A2"
+            description = "Identical to Scenario A1, but with custom probabilities p provided, favoring larger values."
+
+        scenarios.append(
+            _Scenario(
+                letter=letter,
+                description=description,
+                n_k_cons_tuples=[
+                    (n, k, 10)
+                    for n in [10, 100, 1000]
+                    for k in [2**i for i in range(1, 9)]  # 2, 4, 8, ..., 256
+                    if k < n
+                ],
+                use_p=use_p,
+            ),
+        )
+
+    for use_p in [False, True]:
+        if not use_p:
+            letter = "B1"
+            description = "Fixed n=1000 & k=100 with varying number of constraints spanning random 1% portions of the n range (uniform sampling)."
+        else:
+            letter = "B2"
+            description = "Identical to Scenario B1, but with custom probabilities p provided, favoring larger values."
+
+        scenarios.append(
+            _Scenario(
+                letter=letter,
+                description=description,
+                n_k_cons_tuples=[
+                    (1000, 100, n_cons)
+                    for n_cons in [2**i for i in range(1, 10)]  # 2, 4, 8, ..., 512
+                ],
+                use_p=use_p,
+            ),
+        )
+
+    # --- benchmark all scenarios -------------------------
+    print("Benchmarking `randint_constrained`...")
+    print()
+    for s in scenarios:
+        if markdown:
+            print(f"## Scenario {s.letter}")
+        else:
+            print(f"Scenario {s.letter}:")
+
+        print()
+        print(s.description)
+        print()
+
+        # --- create headers --------------------
+        if markdown:
+            headers = [
+                "`k`",
+                "`n`",
+                "`n_cons`",
+                md_multiline(["`randint_numba`", "(time)"]),
+                md_multiline(["`randint_numba`", "(accuracy %)"]),
+                md_multiline(["`randint_constrained_numba`", "(time)"]),
+                md_multiline(["`randint_constrained_numba`", "(accuracy %)"]),
+            ]
+        else:
+            headers = [
+                "k",
+                "n",
+                "n_cons",
+                "randint_numba (time)",
+                "randint_numba (accuracy %)",
+                "randint_constrained_numba (time)",
+                "randint_constrained_numba (accuracy %)",
+            ]
+
+        # --- benchmark scenario ----------------
+        data: list[list[CellContent]] = []
+        for n, k, n_cons in tqdm(s.n_k_cons_tuples, leave=False):
+            # --- build constraints ---
+            if s.letter.startswith("A"):
+                cons = [
+                    Constraint(
+                        int_set=set(range(i * (n // 10), (i + 1) * (n // 10))),
+                        min_count=math.floor(k / 11),
+                        max_count=math.ceil(k / 9),
+                    )
+                    for i in range(10)
+                ]
+            else:
+                cons = []
+                for i in range(n_cons):
+                    cons.append(
+                        Constraint(
+                            int_set=set(
+                                randint_numba(
+                                    n=np.int32(n),
+                                    k=np.int32(n // 100),  # 1% random samples from n
+                                    replace=False,
+                                    seed=np.int64(42 + i),
+                                )
+                            ),
+                            min_count=math.floor(10 / n_cons),
+                            max_count=math.ceil(100 / n_cons),
+                        )
+                    )
+
+            if s.use_p:
+                p = np.array([1.0 + i for i in range(n)], dtype=np.float32)
+                p /= p.sum()
+            else:
+                p = None
+
+            # --- convert constraints to numpy format once ---
+            con_values, con_indices = _build_array_repr(cons)
+
+            # --- benchmark & determine precision ---
+            data.append(
+                [
+                    str(k),
+                    str(n),
+                    str(n_cons),
+                    _benchmark(n, k, con_values, con_indices, p, True, speed),
+                    _determine_precision(n, k, cons, con_values, con_indices, p, True, speed),
+                    _benchmark(n, k, con_values, con_indices, p, False, speed),
+                    _determine_precision(n, k, cons, con_values, con_indices, p, False, speed),
+                ]
+            )
+
+        # --- show results -----------------------------------------
+        data = extend_table_with_aggregate_row(data, agg="geomean", include_percentage=False)
+        data = extend_table_with_aggregate_row(data, agg="mean", include_benchmark_result=False)
+        if markdown:
+            display_data = format_as_markdown(
+                headers,
+                data,
+                highlighters=[
+                    FastestBenchmark(),
+                    HighestPercentage(),
+                    BoldLabels(),
+                ],
+            )
+        else:
+            display_data = format_for_console(headers, data)
+
+        for line in display_data:
+            print(line)
+        print()
+
+
+# =================================================================================================
+#  Internal helpers
+# =================================================================================================
+def _benchmark(
+    n: int,
+    k: int,
+    con_values: np.ndarray[np.int32],
+    con_indices: np.ndarray[np.int32],
+    p: np.ndarray | None,
+    ignore_constraints: bool,
+    speed: float,
+) -> BenchmarkResult:
+    """
+    Runs a benchmark and returns the BenchmarkResult.
+    If ignore_constraints=True, benchmarks randint_numba.
+    If ignore_constraints=False, benchmarks randint_constrained_numba.
+    """
+    n = np.int32(n)
+    k = np.int32(k)
+
+    if ignore_constraints:
+        # Benchmark randint_numba
+        if p is None:
+
+            def benchmark_func():
+                return randint_numba(n=n, k=k, replace=False)
+        else:
+            p_float32 = p.astype(np.float32)
+
+            def benchmark_func():
+                return randint_numba(n=n, k=k, replace=False, p=p_float32)
+    else:
+        # Benchmark randint_constrained_numba
+        if p is None:
+
+            def benchmark_func():
+                return randint_constrained_numba(
+                    n=n,
+                    k=k,
+                    con_values=con_values,
+                    con_indices=con_indices,
+                    p=np.zeros(0, dtype=np.float32),
+                    seed=np.int64(0),
+                )
+        else:
+            p_float32 = p.astype(np.float32)
+
+            def benchmark_func():
+                return randint_constrained_numba(
+                    n=n,
+                    k=k,
+                    con_values=con_values,
+                    con_indices=con_indices,
+                    p=p_float32,
+                    seed=np.int64(0),
+                )
+
+    return benchmark(
+        f=benchmark_func,
+        t_per_run=0.05 / (1000.0**speed),
+        n_warmup=int(8 - 5 * speed),
+        n_benchmark=int(25 - 22 * speed),
+        silent=True,
+    )
+
+
+def _determine_precision(
+    n: int,
+    k: int,
+    cons: list[Constraint],
+    con_values: np.ndarray[np.int32],
+    con_indices: np.ndarray[np.int32],
+    p: np.ndarray | None,
+    ignore_constraints: bool,
+    speed: float,
+) -> Percentage:
+    """
+    Determines how often (%) the constraints are satisfied when sampling.
+    If ignore_constraints=True, samples with randint_numba.
+    If ignore_constraints=False, samples with randint_constrained_numba.
+    """
+
+    # Calculate number of runs based on speed (1000 at speed=0, 2 at speed=1)
+    n_runs = int(1000 * (0.002**speed))
+
+    satisfied_count = 0
+    for run_idx in range(n_runs):
+        # Run the appropriate function with seed equal to run index
+        if ignore_constraints:
+            # Use randint_numba
+            if p is None:
+                result = randint_numba(n=np.int32(n), k=np.int32(k), replace=False, seed=np.int64(run_idx))
+            else:
+                result = randint_numba(
+                    n=np.int32(n), k=np.int32(k), replace=False, p=p.astype(np.float32), seed=np.int64(run_idx)
+                )
+        else:
+            # Use randint_constrained_numba
+            if p is None:
+                result = randint_constrained_numba(
+                    n=np.int32(n),
+                    k=np.int32(k),
+                    con_values=con_values,
+                    con_indices=con_indices,
+                    p=np.zeros(0, dtype=np.float32),
+                    seed=np.int64(run_idx),
+                )
+            else:
+                result = randint_constrained_numba(
+                    n=np.int32(n),
+                    k=np.int32(k),
+                    con_values=con_values,
+                    con_indices=con_indices,
+                    p=p.astype(np.float32),
+                    seed=np.int64(run_idx),
+                )
+
+        # Check if all constraints are satisfied
+        constraints_satisfied = True
+        for constraint in cons:
+            count = sum(1 for val in result if val in constraint.int_set)
+            if count < constraint.min_count or count > constraint.max_count:
+                constraints_satisfied = False
+                break
+
+        if constraints_satisfied:
+            satisfied_count += 1
+
+    return Percentage(frac=satisfied_count / n_runs, decimals=1)

max_div/constraints/__init__.py

@@ -0,0 +1,2 @@
+from .constraint import Constraint
+from .constraints import Constraints

max_div/constraints/_numba.py

@@ -0,0 +1,110 @@
+import numba
+import numpy as np
+from numba.typed import List
+
+from .constraint import Constraint
+
+
+# =================================================================================================
+#  Handlers for numpy-based constraint representation
+# =================================================================================================
+#
+#   Constraints:
+#      [
+#           Constraint(int_set={0,1,2,3,4},   min_count=2, max_count=3),
+#           Constraint(int_set={10,11,12,13}, min_count=0, max_count=7),
+#           Constraint(int_set={3,11},        min_count=2, max_count=2),
+#      ]
+#
+#   Will be represented to 2 numpy arrays:
+#
+#     con_values:
+#         np.array([
+#             [2, 3],      # min_count, max_count for constraint 0
+#             [0, 7],      # min_count, max_count for constraint 1
+#             [2, 2],      # min_count, max_count for constraint 2
+#         ], dtype=np.int32)
+#
+#     con_indices:
+#         -> Part 1 - first 2*n_cons values indicate start/end indices in the array for each constraint
+#         -> Part 2 - followed by concatenated indices from each constraint's int_set
+#
+#                  |-------- Part 1 ----------|----------- Part 2 ---------------|
+#           index:  0   1     2   3     4  5    6      10    11       14    15 16
+#
+#         np.array([6, 11,   11, 15,   15,17,   0,1,2,3,4,   10,11,12,13,    3,11], dtype=np.int32)
+#
+#                     |        |         |      ^^^^^^^^^    ^^^^^^^^^^^     ^^^^
+#                     |        |         |          |            |            |
+#                     |        |         +----->    |            |        con 2 indices
+#                     |        +------------->      |        con 1 indices
+#                     +-------------------->    con 2 indices
+#
+# =================================================================================================
+def _build_array_repr(
+    cons: list[Constraint],
+) -> tuple[np.ndarray[np.int32], np.ndarray[np.int32]]:
+    """
+    Convert list of Constraint objects to numba-compatible representation:
+      - con_values: 2D numpy array of shape (n_cons, 2) with min_count and max_count for each constraint
+      - con_indices: 1D numpy array of shape (total_indices,) with concatenated
+
+    :param cons: list of Constraint objects
+    :return: tuple of (con_values, con_indices)
+    """
+
+    # get dimensions
+    n_cons = len(cons)
+    n_indices = sum([len(con.int_set) for con in cons])
+
+    # pre-allocate
+    con_values = np.empty((n_cons, 2), dtype=np.int32)
+    con_indices = np.empty((2 * n_cons) + n_indices, dtype=np.int32)
+
+    # build con_values
+    for i, con in enumerate(cons):
+        con_values[i, 0] = np.int32(con.min_count)
+        con_values[i, 1] = np.int32(con.max_count)
+
+    # build con_indices
+    i_start = 2 * n_cons  # where we start filling in values from int_set for each constraint
+    for i, con in enumerate(cons):
+        con_indices[2 * i] = np.int32(i_start)
+        con_indices[(2 * i) + 1] = np.int32(i_start + len(con.int_set))
+        for idx in sorted(con.int_set):
+            con_indices[i_start] = np.int32(idx)
+            i_start += 1
+
+    return con_values, con_indices
+
+
+@numba.njit(numba.int32(numba.int32[:, :], numba.int32), inline="always")
+def _np_con_min_value(con_values: np.ndarray[np.int32], i_con: np.int32) -> np.int32:
+    """Return min_value of i-th constraint from con_values array."""
+    return con_values[i_con, 0]
+
+
+@numba.njit(numba.int32(numba.int32[:, :], numba.int32), inline="always")
+def _np_con_max_value(con_values: np.ndarray[np.int32], i_con: np.int32) -> np.int32:
+    """Return max_value of i-th constraint from con_values array."""
+    return con_values[i_con, 1]
+
+
+@numba.njit(numba.int32[:](numba.int32[:], numba.int32), inline="always")
+def _np_con_indices(con_indices: np.ndarray[np.int32], i_con: np.int32) -> np.ndarray[np.int32]:
+    """Return the indices array for the i-th constraint from con_indices array."""
+    start = con_indices[2 * i_con]
+    end = con_indices[2 * i_con + 1]
+    return con_indices[start:end]
+
+
+@numba.njit(inline="always")
+def _np_con_build_index_sets(
+    con_indices: np.ndarray[np.int32],
+    n_cons: np.int32,
+) -> List[set[np.int32]]:
+    """Build list of sets of indices for each constraint from con_indices array."""
+    list_of_sets = List()
+    for i in np.arange(n_cons, dtype=np.int32):
+        list_of_sets.append(set(_np_con_indices(con_indices, i)))
+    return list_of_sets

max_div/constraints/constraint.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Constraint:
+    """Constraint indicating we want to sample at least `min_count` and at most `max_count` integers from `int_set`."""
+
+    int_set: set[int]
+    min_count: int
+    max_count: int

max_div/constraints/constraints.py

@@ -0,0 +1,47 @@
+import copy
+
+import numpy as np
+from numba.typed import List
+
+from .constraint import Constraint
+
+
+class Constraints:
+    """Representation for a collection of constraints."""
+
+    # -------------------------------------------------------------------------
+    #  Constructor / initialization
+    # -------------------------------------------------------------------------
+    def __init__(self):
+        self._cons: list[Constraint] = []
+
+    # -------------------------------------------------------------------------
+    #  API
+    # -------------------------------------------------------------------------
+    @property
+    def n_cons(self) -> int:
+        return len(self._cons)
+
+    def add(self, indices: set[int], min_count: int, max_count: int) -> None:
+        """
+        Add a new constraint, indicating we want to sample at least `min_count` and at most `max_count`
+        integers from `indices`.
+        """
+        self._cons.append(Constraint(indices, min_count, max_count))
+
+    def all(self, deepcopy: bool = False) -> list[Constraint]:
+        """
+        Return list of Constraint objects.
+        :param deepcopy: If True, return a deep copy of the list and its contents.
+        :return: list of constraints representing what was added using `add()`.
+        """
+        if not deepcopy:
+            return self._cons
+        else:
+            return copy.deepcopy(self._cons)
+
+    def to_numpy(self) -> tuple[np.ndarray[np.int32], np.ndarray[np.int32]]:
+        """Convert to 2 numpy arrays (con_values, con_indices) for use in numba sampling functions."""
+        from ._numba import _build_array_repr
+
+        return _build_array_repr(self._cons)

max_div/internal/benchmarking/_micro_benchmark.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import Callable
+from typing import Callable, Literal
 
 import numpy as np
 
@@ -29,6 +31,44 @@ class BenchmarkResult:
         s_perc = f"{50 * (self.t_sec_q_75 - self.t_sec_q_25) / self.t_sec_q_50:.1f}%"
         return f"{s_median} ± {s_perc}"
 
+    def __str__(self) -> str:
+        return self.t_sec_with_uncertainty_str
+
+    @classmethod
+    def aggregate(cls, results: list[BenchmarkResult], method: Literal["mean", "geomean", "sum"]) -> BenchmarkResult:
+        """
+        Aggregate multiple BenchmarkResult objects into a single result, by aggregating q25, q50, 75 values separately.
+
+        :param results: List of BenchmarkResult objects to aggregate
+        :param method: Aggregation method - "mean", "geomean" (geometric mean), or "sum"
+        :return: Aggregated BenchmarkResult
+        """
+        if not results:
+            raise ValueError("Cannot aggregate empty list of results")
+
+        # Collect all quantile values
+        q25_values = [r.t_sec_q_25 for r in results]
+        q50_values = [r.t_sec_q_50 for r in results]
+        q75_values = [r.t_sec_q_75 for r in results]
+
+        # Apply the aggregation method
+        if method == "mean":
+            agg_q25 = np.mean(q25_values)
+            agg_q50 = np.mean(q50_values)
+            agg_q75 = np.mean(q75_values)
+        elif method == "geomean":
+            agg_q25 = np.exp(np.mean(np.log(q25_values)))
+            agg_q50 = np.exp(np.mean(np.log(q50_values)))
+            agg_q75 = np.exp(np.mean(np.log(q75_values)))
+        elif method == "sum":
+            agg_q25 = np.sum(q25_values)
+            agg_q50 = np.sum(q50_values)
+            agg_q75 = np.sum(q75_values)
+        else:
+            raise ValueError(f"Unknown aggregation method: {method}")
+
+        return BenchmarkResult(t_sec_q_25=agg_q25, t_sec_q_50=agg_q50, t_sec_q_75=agg_q75)
+
 
 # =================================================================================================
 #  Main benchmarking function
@@ -36,7 +76,7 @@ class BenchmarkResult:
 def benchmark(
     f: Callable,
     t_per_run: float = 0.1,
-    n_warmup: int = 10,
+    n_warmup: int = 5,
     n_benchmark: int = 30,
     silent: bool = False,
 ) -> BenchmarkResult:
@@ -46,7 +86,7 @@ def benchmark(
     :param f: (Callable) Function to benchmark. Should take no arguments.
     :param t_per_run: (float, default=0.1) time in seconds we want to target per benchmarking run.
                       # of executions/run is adjusted to meet this target.
-    :param n_warmup: (int, default=10) Number of warmup runs to perform before benchmarking.
+    :param n_warmup: (int, default=5) Number of warmup runs to perform before benchmarking.
     :param n_benchmark: (int, default=30) Number of benchmark runs to perform.
     :param silent: (bool, default=False) If True, suppresses any output during benchmarking.
     :return: Median estimate of duration/execution of `f` in seconds.
@@ -85,13 +125,14 @@ def benchmark(
             if not silent:
                 print("w", end="")
 
-        # adjust n_executions
+        # adjust n_executions to bring t_tot closer to t_per_run
+        # NOTE: during warmup we adjust n_executions at a log-scale to reach t_per_run target at end of warmup
         t_tot = timer_tot.t_elapsed_sec()
         n_executions = round(
             clip(
-                value=n_executions * (t_per_run / t_tot),
-                min_value=max(1.0, n_executions / 10),
-                max_value=n_executions * 10,
+                value=n_executions * (t_per_run / t_tot) ** min(1.0, (i + 1) / n_warmup),
+                min_value=max(1.0, n_executions / 100),
+                max_value=n_executions * 100,
             )
         )
 

max_div/internal/formatting/__init__.py

@@ -1 +1,2 @@
+from ._markdown import md_bold, md_colored, md_italic, md_multiline, md_table
 from ._time_duration import format_long_time_duration, format_short_time_duration, format_time_duration

max_div/internal/formatting/_markdown.py

@@ -0,0 +1,43 @@
+def md_table(data: list[list[str]]) -> list[str]:
+    """Creates a Markdown table from a 2D list of strings, where the first list represents the table headers."""
+
+    # Calculate column widths based on all rows
+    num_cols = len(data[0])
+    col_widths = [max(1, max(len(row[i]) for row in data)) for i in range(num_cols)]  # max width of content and >=1
+
+    result = []
+
+    # Add header row
+    header = "| " + " | ".join(data[0][i].ljust(col_widths[i]) for i in range(num_cols)) + " |"
+    result.append(header)
+
+    # Add separator row
+    separator = "| " + " | ".join("-" * col_widths[i] for i in range(num_cols)) + " |"
+    result.append(separator)
+
+    # Add data rows
+    for row in data[1:]:
+        row_str = "| " + " | ".join(row[i].ljust(col_widths[i]) for i in range(num_cols)) + " |"
+        result.append(row_str)
+
+    return result
+
+
+def md_multiline(lines: list[str]) -> str:
+    """Puts multiple Markdown lines into a single line by means of html line breaks."""
+    return "<br>".join(lines)
+
+
+def md_bold(text: str) -> str:
+    """Makes the given text bold in Markdown."""
+    return f"**{text}**"
+
+
+def md_italic(text: str) -> str:
+    """Makes the given text italic in Markdown."""
+    return f"*{text}*"
+
+
+def md_colored(text: str, hex_color: str) -> str:
+    """Colors the given text in Markdown using HTML span tags (hex_color="#rrggbb" or "#rgb")."""
+    return f'<span style="color:{hex_color}">{text}</span>'

max_div/internal/math/__init__.py

@@ -0,0 +1 @@
+