feat: measure the shortlist funnel loss, then size the shortlist per patient

scripts/compare_shortlist_ab.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python
+"""Compare the arms of the shortlist-policy A/B (scripts/run_trec_shortlist_ab.slurm).
+
+Reports each arm against the fixed baseline, and -- crucially -- separates the two
+effects that a depth policy produces:
+
+  * spending compute BETTER  (equal-cost arm: same mean shortlist size, more recall)
+  * spending compute MORE    (deeper arm: bigger shortlist, bought with GPU time)
+
+Conflating them overstates the result, so cost is always printed beside quality.
+
+    uv run python scripts/compare_shortlist_ab.py [--root shortlist_ab] [--track 21]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+BASELINE = "fixed"
+# Ordered deliberately. ndcg_full@10 normalizes by the ideal over the FULL judged pool, so a
+# relevant trial that never entered the shortlist counts as a miss -- it is recall-aware, and
+# it is the metric a depth change should be judged on.
+QUALITY = (
+    "shortlist_recall",
+    "ndcg_full@10",
+    "P@10(eligible)",
+    "P@10(rel>=1)",
+    "recall@1000",
+)
+COST = ("shortlist_size",)
+# ndcg@10 sits here, NOT in QUALITY. It normalizes by the ideal over judged-AND-RANKED trials,
+# which makes it recall-INDEPENDENT by construction: ranking fewer trials shrinks the ideal
+# too. Replaying finished runs at reduced depth shows it moving the wrong way -- on TREC 2021
+# it reads 0.8948 at depth 10 and 0.8215 at depth 196, while ndcg_full@10 goes 0.6028 -> 0.8197
+# and P@10(eligible) 0.3853 -> 0.7680 over the same range. Judging a depth policy on ndcg@10
+# would conclude that more depth hurts, which is exactly backwards.
+DIAGNOSTIC = ("ndcg@10", "funnel_depth_loss", "shortlist_selection_delta")
+
+
+def load_arm(root: Path, arm: str, track: str) -> dict | None:
+    path = root / arm / f"results_trec{track}" / "evaluation_metrics.json"
+    if not path.exists():
+        return None
+    data = json.loads(path.read_text())
+    return {"mean": data.get("mean", {}), "n": data.get("num_queries_scored")}
+
+
+def fmt(value: object, width: int = 9) -> str:
+    if isinstance(value, (int, float)):
+        return f"{value:{width}.4f}"
+    return f"{'--':>{width}}"
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--root", default="shortlist_ab")
+    parser.add_argument("--track", default="21")
+    args = parser.parse_args()
+
+    root = Path(args.root)
+    arms = sorted(p.name for p in root.iterdir() if p.is_dir()) if root.is_dir() else []
+    if BASELINE not in arms:
+        print(f"No '{BASELINE}' arm under {root}/ -- nothing to compare against.")
+        return 1
+
+    loaded = {arm: load_arm(root, arm, args.track) for arm in arms}
+    missing = [arm for arm, data in loaded.items() if data is None]
+    for arm in missing:
+        print(f"NOTE: arm '{arm}' has no evaluation_metrics.json yet (still running?)")
+        loaded.pop(arm)
+    if BASELINE not in loaded:
+        return 1
+
+    base = loaded[BASELINE]["mean"]
+    order = [BASELINE] + [a for a in loaded if a != BASELINE]
+
+    print(f"\nTREC {args.track} shortlist-policy A/B  (n={loaded[BASELINE]['n']} topics)\n")
+    label_w = max(len(a) for a in order) + 2
+    for group, keys in (("COST", COST), ("QUALITY", QUALITY), ("DIAGNOSTIC", DIAGNOSTIC)):
+        print(f"  {group}")
+        print(f"    {'arm':<{label_w}}" + "".join(f"{k:>26s}" for k in keys))
+        for arm in order:
+            mean = loaded[arm]["mean"]
+            cells = ""
+            for key in keys:
+                value = mean.get(key)
+                cell = fmt(value)
+                if arm != BASELINE and isinstance(value, (int, float)):
+                    ref = base.get(key)
+                    if isinstance(ref, (int, float)):
+                        cell += f" ({value - ref:+.4f})"
+                cells += f"{cell:>26s}"
+            print(f"    {arm:<{label_w}}" + cells)
+        print()
+
+    print("  Reading this:")
+    print("    Judge a depth change on ndcg_full@10 and P@10(eligible). Both are recall-aware.")
+    print("    Do NOT judge it on ndcg@10 -- that is normalized over judged-AND-ranked trials,")
+    print("    so it is recall-independent and moves the WRONG way as depth grows.")
+    print("    An arm at the SAME shortlist_size as 'fixed' with higher recall-aware quality")
+    print("    spends its compute better -- that gain is free.")
+    print("    An arm with a LARGER shortlist_size bought its gain with GPU time; compare")
+    print("    it against 'fixed' only after noting the extra cost.")
+    print("    shortlist_selection_delta < 0 means the second level is still selecting worse")
+    print("    than a plain first-level cut at that depth, independent of the policy.\n")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/trialmatchai/config/settings.py

@@ -174,6 +174,22 @@ class FirstLevelSearchSettings(BaseModel):
     )
 
 
+class ShortlistSettings(BaseModel):
+    """How deep the shortlist handed to the eligibility reasoner goes.
+
+    "fixed" keeps the divisor-based sizing (one depth for every patient) and is the
+    default. "relative_to_max" sizes each patient from its own first-level score curve,
+    keeping trials scoring at least ``relative_to_max_alpha`` x that patient's top score.
+    See matching/shortlist_depth.py for the offline evidence.
+    """
+
+    policy: Literal["fixed", "relative_to_max"] = "fixed"
+    relative_to_max_alpha: float = Field(0.25, gt=0.0, le=1.0)
+    min_depth: int = Field(50, ge=1)
+    # None -> bounded only by what the reasoner can consume (rag.max_trials_rag).
+    max_depth: int | None = Field(None, ge=1)
+
+
 class SearchSettings(BaseModel):
     mode: Literal["bm25", "vector", "hybrid"] = "hybrid"
     vector_score_threshold: float = Field(0.5, ge=0.0, le=1.0)
@@ -191,6 +207,7 @@ class SearchSettings(BaseModel):
     first_level: FirstLevelSearchSettings = Field(
         default_factory=FirstLevelSearchSettings
     )
+    shortlist: ShortlistSettings = Field(default_factory=ShortlistSettings)
 
     @model_validator(mode="before")
     @classmethod

src/trialmatchai/main.py

@@ -15,6 +15,8 @@
     rank_trials,
     save_ranked_trials,
 )
+from trialmatchai.matching.query_expansion import build_first_level_expander
+from trialmatchai.matching.shortlist_depth import choose_shortlist_depth, depth_report
 from trialmatchai.matching.retrieval.trial_retrieval import ClinicalTrialSearch
 from trialmatchai.matching.retrieval.criteria_retrieval import SecondStageRetriever
 from trialmatchai.matching.retrieval.location import (
@@ -107,6 +109,7 @@ def run_first_level_search(
     config: Dict,
     search_backend,
     patient_profile: PatientProfile | None = None,
+    llm_query_expander=None,
 ) -> Optional[Tuple]:
     main_conditions = list(keywords.get("main_conditions", []))
     other_conditions = list(keywords.get("other_conditions", []))
@@ -125,6 +128,9 @@ def run_first_level_search(
         search_backend=search_backend,
         embedder=embedder,
         entity_annotator=entity_annotator,
+        # Without this the llm_expansion channel is dead: the planner logs "no expander is
+        # configured" and returns [], however the config flag is set.
+        llm_query_expander=llm_query_expander,
     )
 
     search_cfg = config["search"]
@@ -306,29 +312,57 @@ def run_second_level_search(
     second_level_scores = {
         trial["nct_id"]: trial["score"] for trial in second_level_results
     }
-
+    # Persist the whole second-level pool, not just the shortlist. Together with
+    # first_level_scores.json this makes shortlist fusion replayable offline, so fusion
+    # weights can be retuned against completed runs instead of costing a GPU job each time.
+    # Measured motivation: shortlist_selection_delta is negative on every run so far, i.e.
+    # the fused shortlist selects worse than a plain first-level cut at the same depth.
+    write_json_file(second_level_scores, f"{output_folder}/second_level_scores.json")
+
+    search_config = config.get("search", {})
     combined_scores = _fuse_shortlist_scores(
         nct_ids=nct_ids,
         second_level_results=second_level_results,
         first_level_scores=first_level_scores,
-        search_config=config.get("search", {}),
+        search_config=search_config,
     )
 
     sorted_trials = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)
-    keep_divisor = max(1, int(config.get("search", {}).get("second_level_keep_divisor", 3)))
+    keep_divisor = max(1, int(search_config.get("second_level_keep_divisor", 3)))
     # Size the shortlist off the reranked count, not the fused pool: rank fusion adds
     # first-level-only trials, so keying the divisor to the pool would silently enlarge the
     # shortlist and confound a fusion A/B.
     reranked_count = len(second_level_results) or len(sorted_trials)
     num_top = max(1, min(reranked_count // keep_divisor, top_n))
     # RAG only reasons over rag.max_trials_rag trials; cap the shortlist to match, else trials
     # past the cap get no eligibility output and are silently dropped from the final ranking.
+    upper_bound = len(sorted_trials) or 1
     if _rag_enabled(config):
-        num_top = max(1, min(num_top, int(config.get("rag", {}).get("max_trials_rag", 20))))
+        upper_bound = min(upper_bound, int(config.get("rag", {}).get("max_trials_rag", 20)))
+        num_top = max(1, min(num_top, upper_bound))
+    # The divisor sizes every patient the same. Depth is 94% of the measured shortlist recall
+    # loss and the depth patients need spans 50-1550 trials, so an opt-in policy may widen or
+    # narrow this per patient from the first-level score curve. Default policy returns num_top.
+    fixed_depth = num_top
+    num_top = choose_shortlist_depth(
+        first_level_scores=first_level_scores,
+        fixed_depth=fixed_depth,
+        upper_bound=upper_bound,
+        search_config=search_config,
+    )
     semi_final_trials = sorted_trials[:num_top]
 
     top_trials_path = f"{output_folder}/top_trials.txt"
     write_text_file([trial_id for trial_id, _ in semi_final_trials], top_trials_path)
+    write_json_file(
+        depth_report(
+            chosen=len(semi_final_trials),
+            fixed_depth=fixed_depth,
+            first_level_scores=first_level_scores,
+            search_config=search_config,
+        ),
+        f"{output_folder}/shortlist_depth.json",
+    )
     constraints_config = config.get("constraints", {})
     if constraints_config.get("enabled", True) and constraints_config.get(
         "write_reports",
@@ -524,6 +558,12 @@ def main_pipeline(
 
     embedder = build_embedder(config)
     entity_annotator = build_entity_annotator(config, embedder=embedder)
+    # Built once for the whole run, not per patient: it shares the cached CoT engine, and
+    # rebuilding per patient would re-resolve that engine 75 times. None when
+    # search.first_level.llm_expansion_enabled is off, which is the default.
+    llm_query_expander = build_first_level_expander(config)
+    if llm_query_expander is not None:
+        logger.info("First-level LLM query expansion is ON (llm_expansion channel active).")
 
     with warnings.catch_warnings():
         warnings.filterwarnings(
@@ -618,6 +658,7 @@ def main_pipeline(
                         config,
                         search_backend,
                         patient_profile=profile,
+                        llm_query_expander=llm_query_expander,
                     )
             if not result:
                 logger.error("First-level search failed for %s", patient_id)