Sweep index size and negative ratio in run_grid_search

Index size is the highest-leverage knob for tuning, but run_grid_search could
only sweep top_k and min_score. With subsample() available, size and ratio can
be swept too, and cheaply: reshaping the index costs about a millisecond, while
re-scoring costs seconds.

The loop is organised around what each setting costs. Size and ratio reshape the
index (cheap) and become an outer loop; top_k forces a re-score (expensive);
thresholds and aggregators stay free on the cached scores.

Rows now carry the requested n_positive and neg_to_pos_ratio plus the actual
n_positive_actual and n_negative_actual counts. The actual counts matter because
a request is clipped when the index is smaller than asked for - without them two
rows can look identical while describing different indices.

Backward compatible: both new arguments default to None, meaning one pass using
the index exactly as given. subsample() is not called at all in that case, so
any index-like object keeps working and existing behaviour is unchanged.

Also fixes a latent NameError: simulation.py referenced LOG without importing
logging or defining it. Nothing had triggered it because the module had no
logging calls until this one.

Documented in the docstring, and worth calling out for reviewers: the cost is
len(sizes) x len(ratios) x len(top_k) scoring passes, and the observation texts
are re-encoded on every pass even though their embeddings do not depend on the
index at all. Letting callers pass pre-computed sample embeddings would collapse
that to roughly one encoding pass. Deliberately out of scope here.

Adds 5 tests covering the default no-subsample path, both axes together and
separately, clipped counts, and the per-configuration progress logging.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: wxiao0421

src/sentinel/simulation.py

@@ -56,6 +56,7 @@
     rows = compare_aggregators(scored)
 """
 
+import logging
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Callable, Dict, List, Mapping, Optional, Sequence
 
@@ -75,6 +76,8 @@
     # stays lightweight (numpy-only) and does not pull in torch / transformers.
     from sentinel.sentinel_local_index import SentinelLocalIndex
 
+LOG = logging.getLogger(__name__)
+
 
 # A convenient name -> function map of the built-in summarize metrics, so a
 # caller (or a notebook) can iterate over "all aggregators" in one line.
@@ -515,6 +518,21 @@ def compare_aggregators(
     ]
 
 
+def _row_count(embeddings: object) -> Optional[int]:
+    """Return the number of rows in an embedding tensor, or None if unavailable.
+
+    Args:
+        embeddings: A tensor-like object, or None.
+
+    Returns:
+        The row count, or None when the object has no usable shape.
+    """
+    shape = getattr(embeddings, "shape", None)
+    if shape is None or len(shape) == 0:
+        return None
+    return int(shape[0])
+
+
 def run_grid_search(
     index: "SentinelLocalIndex",
     groups: Sequence[LabeledGroup],
@@ -525,13 +543,33 @@ def run_grid_search(
     top_n: Optional[int] = None,
     decision_threshold: Optional[float] = None,
     show_progress_bar: bool = False,
+    n_positive_values: Optional[Sequence[int]] = None,
+    neg_to_pos_ratios: Optional[Sequence[float]] = None,
+    index_seed: Optional[int] = None,
 ) -> List[Dict[str, float]]:
     """Sweep hyperparameters and summarize metrics, returning a flat result table.
 
-    For each ``top_k`` the groups are scored once (the expensive step), then
-    every combination of ``min_score_values`` x ``aggregators`` is evaluated
-    cheaply. The returned rows are easy to turn into a ``pandas.DataFrame`` for
-    plotting.
+    The sweep is organised around what each setting costs. Index size and ratio
+    reshape the index (cheap, via
+    :meth:`~sentinel.sentinel_local_index.SentinelLocalIndex.subsample`), ``top_k``
+    forces a re-scoring (expensive), and thresholds and aggregators are evaluated
+    for free on the cached scores::
+
+        for n_positive in n_positive_values:      # cheap: subsample()
+          for ratio in neg_to_pos_ratios:         # cheap: subsample()
+            for top_k in top_k_values:            # EXPENSIVE: re-scores
+              for min_score in min_score_values:  # cheap
+                for aggregator in aggregators:    # cheap
+
+    The returned rows are easy to turn into a ``pandas.DataFrame`` for plotting.
+
+    Cost warning: the number of scoring passes is
+    ``len(n_positive_values) x len(neg_to_pos_ratios) x len(top_k_values)``. A 7x7
+    size/ratio grid with three ``top_k`` values is 147 full passes. Note that the
+    observation texts are re-encoded on every pass even though their embeddings do
+    not depend on the index at all, so most of that cost is recomputing identical
+    numbers. Letting callers pass pre-computed sample embeddings would collapse it
+    to roughly one encoding pass; that is a separate change.
 
     Args:
         index: A loaded :class:`~sentinel.sentinel_local_index.SentinelLocalIndex`.
@@ -543,30 +581,94 @@ def run_grid_search(
         decision_threshold: Cutoff for the classification family (``None`` =
             best-F1).
         show_progress_bar: Whether to show the encoder progress bar.
+        n_positive_values: Index sizes to try, as counts of positive examples.
+            ``None`` (the default) means one pass using the index exactly as given.
+        neg_to_pos_ratios: Negative-to-positive ratios to try. ``None`` (the
+            default) means one pass using the index exactly as given.
+        index_seed: Seed for the subsampling, so each index configuration is
+            reproducible across runs.
 
     Returns:
-        A list of metric dicts, one per ``(top_k, min_score_to_consider,
-        aggregator)`` combination. Each row also includes a ``top_k`` key.
+        A list of metric dicts, one per ``(n_positive, neg_to_pos_ratio, top_k,
+        min_score_to_consider, aggregator)`` combination. Each row also includes
+        ``top_k``, the requested ``n_positive`` and ``neg_to_pos_ratio``, and the
+        actual ``n_positive_actual`` / ``n_negative_actual`` counts. The actual
+        counts matter because a request is clipped when the index is smaller than
+        asked for - without them two rows can look identical while describing
+        different indices.
     """
     if aggregators is None:
         aggregators = DEFAULT_AGGREGATORS
 
+    # None means "one pass, use the index exactly as given", which reproduces the
+    # behaviour from before these axes existed.
+    positive_options: Sequence[Optional[int]] = (
+        list(n_positive_values) if n_positive_values is not None else [None]
+    )
+    ratio_options: Sequence[Optional[float]] = (
+        list(neg_to_pos_ratios) if neg_to_pos_ratios is not None else [None]
+    )
+    total_configurations = len(positive_options) * len(ratio_options)
+
     rows: List[Dict[str, float]] = []
-    for top_k in top_k_values:
-        scored = score_groups(
-            index, groups, top_k=top_k, show_progress_bar=show_progress_bar
-        )
-        for min_score in min_score_values:
-            for name, fn in aggregators.items():
-                row = evaluate_groups(
-                    scored,
-                    fn,
-                    aggregator_name=name,
-                    min_score_to_consider=min_score,
-                    top_n=top_n,
-                    decision_threshold=decision_threshold,
+    configuration = 0
+    for n_positive in positive_options:
+        for ratio in ratio_options:
+            configuration += 1
+
+            if n_positive is None and ratio is None:
+                # Never call subsample() in the default case, so this keeps working
+                # with any index-like object and stays byte-identical to before.
+                working_index = index
+            else:
+                working_index = index.subsample(
+                    n_positive=n_positive,
+                    neg_to_pos_ratio=ratio,
+                    seed=index_seed,
+                )
+
+            n_positive_actual = _row_count(
+                getattr(working_index, "positive_embeddings", None)
+            )
+            n_negative_actual = _row_count(
+                getattr(working_index, "negative_embeddings", None)
+            )
+
+            if total_configurations > 1:
+                # A long sweep is otherwise indistinguishable from a hang.
+                LOG.info(
+                    "Grid search index configuration %d/%d: n_positive=%s ratio=%s "
+                    "(actual %s positives / %s negatives)",
+                    configuration,
+                    total_configurations,
+                    n_positive,
+                    ratio,
+                    n_positive_actual,
+                    n_negative_actual,
+                )
+
+            for top_k in top_k_values:
+                scored = score_groups(
+                    working_index,
+                    groups,
+                    top_k=top_k,
+                    show_progress_bar=show_progress_bar,
                 )
-                row["top_k"] = int(top_k)
-                rows.append(row)
+                for min_score in min_score_values:
+                    for name, fn in aggregators.items():
+                        row = evaluate_groups(
+                            scored,
+                            fn,
+                            aggregator_name=name,
+                            min_score_to_consider=min_score,
+                            top_n=top_n,
+                            decision_threshold=decision_threshold,
+                        )
+                        row["top_k"] = int(top_k)
+                        row["n_positive"] = n_positive
+                        row["neg_to_pos_ratio"] = ratio
+                        row["n_positive_actual"] = n_positive_actual
+                        row["n_negative_actual"] = n_negative_actual
+                        rows.append(row)
 
     return rows

tests/test_simulation.py

@@ -19,6 +19,7 @@
 that mimics ``calculate_rare_class_affinity``.
 """
 
+import logging
 from types import SimpleNamespace
 
 import numpy as np
@@ -259,3 +260,135 @@ def test_run_grid_search_rescoring_and_rows():
     # Re-scoring happens once per top_k value (each scores both groups), and NOT
     # again for every threshold/aggregator combination: 2 top_k x 2 groups = 4.
     assert len(index.calls) == 2 * len(groups)
+
+
+class _StubSubsamplableIndex(_StubIndex):
+    """Stub that also records subsample() calls and reports embedding row counts."""
+
+    def __init__(self, n_positive=100, n_negative=500):
+        super().__init__()
+        self.subsample_calls = []
+        self.positive_embeddings = np.zeros((n_positive, 4))
+        self.negative_embeddings = np.zeros((n_negative, 4))
+
+    def subsample(self, n_positive=None, neg_to_pos_ratio=None, seed=None):
+        self.subsample_calls.append(
+            {"n_positive": n_positive, "neg_to_pos_ratio": neg_to_pos_ratio, "seed": seed}
+        )
+        available_positive = self.positive_embeddings.shape[0]
+        kept_positive = (
+            min(n_positive, available_positive) if n_positive else available_positive
+        )
+        kept_negative = (
+            min(int(kept_positive * neg_to_pos_ratio), self.negative_embeddings.shape[0])
+            if neg_to_pos_ratio
+            else self.negative_embeddings.shape[0]
+        )
+        smaller = _StubSubsamplableIndex(kept_positive, kept_negative)
+        # Share the call log so assertions can see scoring done via the copy.
+        smaller.calls = self.calls
+        return smaller
+
+
+def _two_groups():
+    return [
+        LabeledGroup(name="pos", label=1, observations=["0.9", "0.8"]),
+        LabeledGroup(name="neg", label=0, observations=["0.2", "0.1"]),
+    ]
+
+
+def test_grid_search_without_index_axes_never_subsamples():
+    """The default path must not touch the index at all.
+
+    This is the backward-compatibility guarantee: callers passing any index-like
+    object keep working, and behaviour is unchanged from before these axes existed.
+    """
+    index = _StubSubsamplableIndex()
+    rows = run_grid_search(index, _two_groups(), top_k_values=[3], min_score_values=[0.1])
+
+    assert index.subsample_calls == []
+    assert len(rows) == len(DEFAULT_AGGREGATORS)
+    # The new columns are still present, so the table shape is consistent.
+    assert all(row["n_positive"] is None for row in rows)
+    assert all(row["neg_to_pos_ratio"] is None for row in rows)
+
+
+def test_grid_search_sweeps_index_size_and_ratio():
+    """Both new axes multiply out, and each configuration is subsampled once."""
+    index = _StubSubsamplableIndex()
+    rows = run_grid_search(
+        index,
+        _two_groups(),
+        top_k_values=[3, 5],
+        min_score_values=[0.0],
+        n_positive_values=[10, 20],
+        neg_to_pos_ratios=[1.0, 2.0],
+        index_seed=42,
+    )
+
+    # 2 sizes x 2 ratios x 2 top_k x 1 threshold x 6 aggregators.
+    assert len(rows) == 2 * 2 * 2 * len(DEFAULT_AGGREGATORS)
+    # subsample() is called once per (size, ratio) pair, NOT once per top_k: the
+    # whole point is that reshaping the index is cheap and re-scoring is not.
+    assert len(index.subsample_calls) == 4
+    assert all(call["seed"] == 42 for call in index.subsample_calls)
+    assert {(c["n_positive"], c["neg_to_pos_ratio"]) for c in index.subsample_calls} == {
+        (10, 1.0), (10, 2.0), (20, 1.0), (20, 2.0)
+    }
+    # Scoring happens once per (size, ratio, top_k), each covering both groups.
+    assert len(index.calls) == 4 * 2 * len(_two_groups())
+
+
+def test_grid_search_reports_requested_and_actual_counts():
+    """Actual counts are emitted, because a request is clipped on a small index.
+
+    Without them, two rows can look identical while describing different indices.
+    """
+    index = _StubSubsamplableIndex(n_positive=15, n_negative=500)
+    rows = run_grid_search(
+        index,
+        _two_groups(),
+        top_k_values=[3],
+        min_score_values=[0.0],
+        n_positive_values=[10, 999],  # 999 exceeds the 15 available
+        neg_to_pos_ratios=[2.0],
+    )
+
+    by_request = {row["n_positive"]: row for row in rows}
+    assert by_request[10]["n_positive_actual"] == 10
+    assert by_request[10]["n_negative_actual"] == 20
+    # Clipped to what the index actually holds, and visible in the row.
+    assert by_request[999]["n_positive_actual"] == 15
+    assert by_request[999]["n_negative_actual"] == 30
+
+
+def test_grid_search_one_axis_at_a_time():
+    """Either axis can be swept on its own."""
+    index = _StubSubsamplableIndex()
+    size_only = run_grid_search(
+        index, _two_groups(), top_k_values=[3], min_score_values=[0.0],
+        n_positive_values=[10, 20],
+    )
+    assert {row["n_positive"] for row in size_only} == {10, 20}
+    assert all(row["neg_to_pos_ratio"] is None for row in size_only)
+
+    index2 = _StubSubsamplableIndex()
+    ratio_only = run_grid_search(
+        index2, _two_groups(), top_k_values=[3], min_score_values=[0.0],
+        neg_to_pos_ratios=[0.5, 1.0],
+    )
+    assert {row["neg_to_pos_ratio"] for row in ratio_only} == {0.5, 1.0}
+    assert all(row["n_positive"] is None for row in ratio_only)
+
+
+def test_grid_search_logs_progress_per_configuration(caplog):
+    """A long sweep logs progress so it cannot be mistaken for a hang."""
+    caplog.set_level(logging.INFO)
+    index = _StubSubsamplableIndex()
+    run_grid_search(
+        index, _two_groups(), top_k_values=[3], min_score_values=[0.0],
+        n_positive_values=[10, 20], neg_to_pos_ratios=[1.0],
+    )
+
+    assert "index configuration 1/2" in caplog.text
+    assert "index configuration 2/2" in caplog.text