Sweep index size and negative ratio in run_grid_search

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