feat: implement the MCTS explanation method

.gitignore

@@ -169,3 +169,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
+
+jobs_testing/
+.claude/

ciao/__init__.py

@@ -1,7 +1,5 @@
 from ciao.explainer.ciao_explainer import CIAOExplainer
-from ciao.explainer.explanation_methods import (
-    make_lookahead_method,
-)
+from ciao.explainer.explanation_methods import make_lookahead_method, make_mcts_method
 from ciao.model.predictor import ModelPredictor
 from ciao.typing import ExplanationMethodFn
 
@@ -11,4 +9,5 @@
     "ExplanationMethodFn",
     "ModelPredictor",
     "make_lookahead_method",
+    "make_mcts_method",
 ]

ciao/__main__.py

@@ -75,14 +75,25 @@ def _log_trajectory(run_id: str, results: ExplanationResult) -> None:
     """Batch-log per-region trajectory points to MLflow."""
     ts_ms = int(time.time() * 1000)
     trajectory_metrics = [
-        Metric(
-            key=f"region_{idx}/trajectory_best_score",
-            value=item["best_score"],
-            timestamp=ts_ms,
-            step=int(item["evals"]),
-        )
+        metric
         for idx, region in enumerate(results.regions)
-        for item in region.trajectory
+        # Deduplicate by evals: multiple iterations may share the same eval_count
+        # (e.g. cached terminal hits in MCGS), which would violate MLflow's metric PK.
+        for item in {item["evals"]: item for item in region.trajectory}.values()
+        for metric in (
+            Metric(
+                key=f"region_{idx}/trajectory_best_score",
+                value=item["best_score"],
+                timestamp=ts_ms,
+                step=int(item["evals"]),
+            ),
+            Metric(
+                key=f"region_{idx}/trajectory_time",
+                value=item["time"],
+                timestamp=ts_ms,
+                step=int(item["evals"]),
+            ),
+        )
     ]
     if not trajectory_metrics:
         return
@@ -96,15 +107,18 @@ def _log_trajectory(run_id: str, results: ExplanationResult) -> None:
 
 
 def _log_explanation_results(
-    run_id: str, results: ExplanationResult, elapsed: float
+    run_id: str,
+    results: ExplanationResult,
+    elapsed: float,
 ) -> None:
-    """Log explanation params, per-region metrics, trajectory, and timing to MLflow."""
+    """Log explanation params, baseline log-odds, per-region metrics, trajectory, and timing to MLflow."""
     mlflow.log_params(
         {
             "target_class_idx": results.target_class_idx,
             "class_name": results.class_name,
         }
     )
+    mlflow.log_metric("original_log_odds", results.original_log_odds)
 
     for idx, region in enumerate(results.regions):
         mlflow.log_metrics(
@@ -207,6 +221,7 @@ def main(cfg: DictConfig) -> None:
                     predictor=predictor,
                     target_class_idx=cfg.target_class_idx,
                     max_regions=cfg.max_regions,
+                    sigma=cfg.sigma,
                     desired_length=cfg.desired_length,
                     batch_size=cfg.batch_size,
                     segmentation=segmentation,

ciao/algorithm/__init__.py

@@ -4,6 +4,7 @@
 from ciao.algorithm.context import SearchContext
 from ciao.algorithm.graph import ImageGraph
 from ciao.algorithm.lookahead import build_region_greedy_lookahead
+from ciao.algorithm.mcts import build_region_mcts
 from ciao.algorithm.search_helpers import is_terminal
 
 
@@ -12,5 +13,6 @@
     "SearchContext",
     "build_all_regions",
     "build_region_greedy_lookahead",
+    "build_region_mcts",
     "is_terminal",
 ]

ciao/algorithm/builder.py

@@ -1,5 +1,7 @@
 """Unified region builder orchestrating different search algorithms."""
 
+from typing import Literal
+
 import torch
 
 from ciao.algorithm.context import SearchContext
@@ -9,15 +11,44 @@
 from ciao.typing import ExplanationMethodFn
 
 
+SeedSelectionMode = Literal[-1, 1] | None
+
+
+def _select_seed_and_sign(
+    scores: dict[int, float], available_segments: list[int], sigma: SeedSelectionMode
+) -> tuple[int, int]:
+    if not available_segments:
+        raise ValueError(
+            "Cannot select a seed from an empty set of available segments."
+        )
+
+    if sigma is None:
+        seed_idx = max(available_segments, key=lambda x: abs(scores[x]))
+        optimization_sign = 1 if scores[seed_idx] >= 0 else -1
+    elif sigma == 1:
+        seed_idx = max(available_segments, key=lambda x: scores[x])
+        optimization_sign = 1
+    elif sigma == -1:
+        seed_idx = max(available_segments, key=lambda x: -scores[x])
+        optimization_sign = -1
+    else:
+        raise ValueError(f"sigma must be -1, 1, or 'auto', got {sigma!r}")
+
+    return seed_idx, optimization_sign
+
+
 def build_all_regions(
     method: ExplanationMethodFn,
     predictor: ModelPredictor,
     input_batch: torch.Tensor,
     replacement_image: torch.Tensor,
     image_graph: ImageGraph,
     target_class_idx: int,
+    original_log_odds: torch.Tensor,
     scores: dict[int, float],
     max_regions: int,
+    original_prob: float,
+    sigma: SeedSelectionMode = None,
     desired_length: int = 30,
     batch_size: int = 64,
 ) -> list[RegionResult]:
@@ -34,8 +65,13 @@ def build_all_regions(
         replacement_image: Replacement tensor
         image_graph: Graph representation of image segments
         target_class_idx: Target class index
+        original_log_odds: Pre-computed unmasked target-class log-odds (scalar tensor)
         scores: Base segment scores
         max_regions: Maximum number of regions to construct
+        original_prob: Pre-computed unmasked probability for the target class
+        sigma: Seed selection mode. ``"None"`` picks max abs score and then inherits its sign. ``1`` chooses the
+            highest raw score (positive evidence), ``-1`` chooses the lowest
+            raw score (negative evidence).
         desired_length: Target number of segments per region
         batch_size: Batch size for model evaluation
 
@@ -45,8 +81,6 @@ def build_all_regions(
     regions: list[RegionResult] = []
     used_segments: set[int] = set()
 
-    original_prob = predictor.get_predictions(input_batch)[0, target_class_idx].item()
-
     for _ in range(max_regions):
         # Find best unprocessed seed
         available_segments = [
@@ -56,9 +90,9 @@ def build_all_regions(
         if not available_segments:
             break
 
-        seed_idx = max(available_segments, key=lambda x: abs(scores[x]))
-        seed_score = scores[seed_idx]
-        optimization_sign = 1 if seed_score >= 0 else -1
+        seed_idx, optimization_sign = _select_seed_and_sign(
+            scores=scores, available_segments=available_segments, sigma=sigma
+        )
 
         # Construct a SearchContext for the current step
         ctx = SearchContext(
@@ -67,6 +101,7 @@ def build_all_regions(
             replacement_image=replacement_image,
             image_graph=image_graph,
             target_class_idx=target_class_idx,
+            original_log_odds=original_log_odds,
             seed_idx=seed_idx,
             optimization_sign=optimization_sign,
             used_segments=frozenset(used_segments),

ciao/algorithm/context.py

@@ -15,6 +15,7 @@ class SearchContext:
     replacement_image: torch.Tensor
     image_graph: ImageGraph
     target_class_idx: int
+    original_log_odds: torch.Tensor
     seed_idx: int
     optimization_sign: int
     used_segments: frozenset[int]

ciao/algorithm/lookahead.py

@@ -3,6 +3,7 @@
 Rolling horizon strategy: Look ahead multiple steps but only commit one step at a time.
 """
 
+import time
 from collections import deque
 from collections.abc import Set
 
@@ -45,6 +46,7 @@ def build_region_greedy_lookahead(
     eval_count = 0
     trajectory: list[dict[str, float]] = []
     current_best_score = -float("inf")
+    t0 = time.monotonic()
 
     # Grow region one step at a time
     while len(current_region) < desired_length:
@@ -71,6 +73,7 @@ def build_region_greedy_lookahead(
             replacement_image=ctx.replacement_image,
             target_class_idx=ctx.target_class_idx,
             batch_size=ctx.batch_size,
+            original_log_odds=ctx.original_log_odds,
         )
         eval_count += len(candidate_regions)
 
@@ -85,7 +88,13 @@ def build_region_greedy_lookahead(
         signed_best = best_score * optimization_sign
         if signed_best > current_best_score:
             current_best_score = signed_best
-        trajectory.append({"evals": eval_count, "best_score": current_best_score})
+        trajectory.append(
+            {
+                "evals": eval_count,
+                "best_score": current_best_score,
+                "time": time.monotonic() - t0,
+            }
+        )
 
         # Optimization - commit entire path
         if len(best_region) == desired_length:
@@ -109,9 +118,16 @@ def build_region_greedy_lookahead(
             replacement_image=ctx.replacement_image,
             target_class_idx=ctx.target_class_idx,
             batch_size=1,
+            original_log_odds=ctx.original_log_odds,
         )[0]
         eval_count += 1
-        trajectory.append({"evals": eval_count, "best_score": current_best_score})
+        trajectory.append(
+            {
+                "evals": eval_count,
+                "best_score": current_best_score,
+                "time": time.monotonic() - t0,
+            }
+        )
 
     return RegionResult(
         region=current_region,