Codex/balanced median refinement

README.md

@@ -133,10 +133,20 @@ The Overlap Index can be used in several settings:
   require dense feature arrays.
 - Sparse feature matrices remain sparse during fitting, slicing, prediction,
   multi-label expansion, and overlap scoring. KMeans centroids and the bounded
-  prototype-distance score blocks are still dense.
+  prototype-distance score blocks are still dense. When balanced-median
+  refinement is enabled, only an eligible parent's local support block is
+  densified while its observation representatives are selected.
 - Offline scoring uses a backend-neutral scratch planner. The default
   `offline_memory_budget_mb=256` is a scratch-memory budget only; it does not
   cap the fitted model or retain a full sample-by-prototype score matrix.
+- KMeans and MiniBatchKMeans optionally support the deterministic one-pass
+  prototype refinement enabled by `prototype_refinement=True`. It is
+  single-label only and replaces eligible (support >= 2, zero outgoing
+  own-runner-up count) parents with balanced, observation-based median
+  representatives; the default `False` leaves fitted centers unchanged.
+  See the backend guides for the runtime/prototype-resolution trade-off and
+  fitted diagnostics. The fitted `prototype_refinement_` summary reports the
+  resolved mode as the strings `"none"` or `"balanced_median"`.
 - Normalize input features before fitting. Examples in this repository use `MinMaxScaler` for convenience.
 - ART backends complement-code inputs internally and therefore require features in the `[0, 1]` interval.
 - Offline backends (`KMeans`, `MiniBatchKMeans`, and `BallCover`) consume normalized features directly and do not apply complement coding.
@@ -567,6 +577,19 @@ compare historical COI values directly with scores from this calibration.
   Scratch-memory budget for backend-neutral offline score blocks. This budget
   controls temporary tiles only; it does not limit fitted data or model size.
 
+- `prototype_refinement` *(bool, default=False)*
+  Optional deterministic one-pass refinement for `KMeans` and
+  `MiniBatchKMeans`. Set it to `True` for scalar single-label fits; it uses
+  the fit-time eligibility rule (support >= 2 and zero outgoing own-runner-up
+  count), projects each selected parent into balanced halves, and chooses an
+  actual observation nearest each coordinate-wise median. `False` preserves
+  the ordinary fitted centers and is accepted for every backend; enabling the
+  option on an unsupported backend raises an error. The option does not fit
+  child estimators or run gates/rescue logic. Sparse inputs remain accepted,
+  although a local support block may be densified. Inspect the fitted
+  `prototype_refinement_` summary (whose resolved mode is `"none"` or
+  `"balanced_median"`); `score_fixed` uses the resulting centers.
+
 - `multilabel_pair_mode` *("all" or "top_m")*
   Directional competitor selection strategy for multi-label offline scoring.
 

docs/backends/kmeans.md

@@ -33,6 +33,40 @@ print(oi.index)
 `n_init` is omitted, the adapter uses `n_init="auto"`. `kmeans_k` may be an
 integer or a dictionary keyed by every observed label.
 
+### Optional balanced-median refinement
+
+Set `prototype_refinement=True` to opt in to a deterministic, one-pass
+balanced-median refinement after the KMeans fit:
+
+```python
+oi = OverlapIndex(
+    model_type="KMeans",
+    kmeans_k=10,
+    kmeans_kwargs={"random_state": 0, "n_init": "auto"},
+    prototype_refinement=True,
+)
+oi.fit(X, y)  # y must contain one scalar label per row
+print(oi.prototype_refinement_["applied_count"])
+```
+
+Eligibility is frozen from the fitted data: a parent prototype must have
+support of at least two and zero outgoing own-runner-up count. Each eligible
+parent is projected along a deterministic farthest-pair axis, split into
+balanced halves, and represented by the actual observation nearest each
+half's coordinate-wise median. No child scikit-learn fit, gate, or rescue
+step is run. The pass is single-label only; multi-label targets are rejected.
+`score_fixed` scores with these already-refined centers and does not refit or
+refine them again.
+
+The default `prototype_refinement=False` preserves the ordinary KMeans
+centers. Refinement can improve resolution for broad, isolated supports, but
+it adds a tiled fit-isolation pass and can increase the prototype count and
+subsequent scoring cost. Sparse input is accepted; only local support blocks
+used by a candidate split may be densified. Use the opt-in when that runtime
+and prototype-growth trade-off is acceptable. The fitted
+`prototype_refinement_` diagnostics expose the resolved mode as
+`"none"` or `"balanced_median"`.
+
 ## Tuning guidance
 
 - Hold `kmeans_k` constant when comparing representations.

docs/backends/minibatch_kmeans.md

@@ -45,6 +45,39 @@ OverlapIndex uses the fitted centers rather than the estimator's `labels_` or
 exact inertia. `kmeans_k` may be one positive integer for every label or a
 dictionary of label-specific counts.
 
+### Optional balanced-median refinement
+
+Pass `prototype_refinement=True` when a single-label fit can benefit from a
+deterministic balanced-median prototype split:
+
+```python
+oi = OverlapIndex(
+    kmeans_k=10,
+    kmeans_kwargs={"random_state": 0, "batch_size": 256, "n_init": 1},
+    prototype_refinement=True,
+)
+oi.fit(X, y)  # y must contain one scalar label per row
+print(oi.prototype_refinement_["prototype_count_after"])
+```
+
+The one-pass eligibility check runs on the original fitted centers. A parent
+needs support of at least two and zero outgoing own-runner-up count. Its rows
+are ordered along a deterministic farthest-pair projection, divided into
+balanced halves, and each child is the actual observation nearest that half's
+coordinate-wise median. No child scikit-learn fit, gate, or rescue step is
+performed, and appended children are not reconsidered in the same pass.
+Multi-label targets are rejected for this mode. `score_fixed` uses the
+already-refined centers without another fit or refinement pass.
+
+`prototype_refinement=False` (the default) retains the standard
+MiniBatchKMeans centers. Balanced refinement adds one tiled fit-isolation scan
+and may increase both the prototype count and scoring time. Sparse matrices
+remain supported, although a candidate's local support block may be
+densified. Treat the option as an explicit runtime-versus-prototype-resolution
+trade-off and enable it only when the extra resolution is useful. The fitted
+`prototype_refinement_` diagnostics expose the resolved mode as `"none"` or
+`"balanced_median"`.
+
 ## Tuning guidance
 
 - Increase `kmeans_k` when a label has multimodal or curved support that a few

docs/diagnostics.md

@@ -19,6 +19,7 @@ maps fitted attributes and common warnings back to useful checks.
 | `unevaluable_pairs_` | Lazy set-like view of selected multi-label pairs with no suitable rows. |
 | `unevaluable_labels_` | Labels with no evaluable selected pair. |
 | `n_features_in_` | Feature count recorded at fit time. |
+| `prototype_refinement_` | Fit-time summary of optional centroid refinement decisions. |
 
 The mappings preserve the historical mapping API. Pairwise mappings are sparse:
 iteration and `dict(...)` expose only materialized non-default entries. Direct
@@ -39,6 +40,40 @@ for (source, competitor), score in worst_pairs:
 For multi-label data, filter non-finite scores before sorting because
 unevaluable pairs are represented by `NaN`.
 
+## Prototype refinement diagnostics
+
+For `KMeans` and `MiniBatchKMeans`, `prototype_refinement=False` (the default)
+leaves the fitted centers unchanged. The opt-in `prototype_refinement=True`
+mode is supported for scalar single-label fits only. It performs one
+deterministic pass over the original fit: a parent is eligible when its
+best-match support is at least two and its outgoing own-runner-up count is
+zero. Eligible rows are split along a farthest-pair projection, and each child
+is an actual observation nearest its half's coordinate-wise median. Children
+are not reconsidered in that pass; there is no child scikit-learn fit, gate, or
+rescue step. `score_fixed` uses the resulting fitted centers and does not
+refit or refine them. The public switch is a strict boolean; when enabled on
+an unsupported backend it raises an error, while `False` is accepted without
+refinement. Diagnostics retain the resolved mode names described below.
+
+The fitted `prototype_refinement_` value is a read-mostly mapping with this
+stable top-level schema:
+
+| Key | Meaning |
+| --- | --- |
+| `method`, `mode`, `splitter`, `split_method` | Resolved mode name (`"none"` or `"balanced_median"`). |
+| `prototype_count_before`, `prototype_count_after` | Prototype count before and after the pass. |
+| `eligible_count`, `attempted_count`, `applied_count`, `skipped_count` | Counts for the frozen eligibility and split decisions. |
+| `eligible_parent_ids`, `applied_parent_ids`, `skipped_parent_ids` | Stable tuples of global parent IDs by outcome. |
+| `records` | Tuple of per-parent dictionaries containing `parent`/`parent_id`/`original_id`, class, support, status, reason, child IDs/supports, and selected observation indices (`selected_observation_indices`, plus the `selected_sample_indices` alias when applied). Applied records also include `new_id`. |
+
+When refinement is enabled, the isolation scan is tiled using the offline
+memory and row settings. Sparse input remains supported, but a candidate's
+local support block may be densified to construct its observation
+representatives. Refinement can add prototype-resolution where a broad
+isolated parent hides structure, at the cost of one extra fit-time scan,
+additional centers, and more subsequent scoring work. Keep the default off
+unless that trade-off is useful for the analysis.
+
 `unevaluable_pairs_` is a lazy set-like diagnostic view on multi-label fits:
 iterate it or use membership testing to enumerate/check every selected
 directional pair whose `pairwise_cardinality` is zero. It retains only class

overlapindex/OverlapIndex.py

@@ -33,6 +33,10 @@ class BaseEstimator:  # type: ignore[no-redef]
     compute_second_best_source_scores,
     iter_target_class_blocks,
 )
+from overlapindex._prototype_refinement import (
+    empty_refinement_summary,
+    refinement_method,
+)
 
 
 def _default_one() -> float:
@@ -437,6 +441,7 @@ def __init__(
         top_m: Optional[int] = None,
         exclude_classes: Optional[Any] = None,
         offline_memory_budget_mb: int = 256,
+        prototype_refinement: bool = False,
     ) -> None:
         """
         Initialize the overlap index and its clustering backend.
@@ -487,6 +492,12 @@ def __init__(
             prototypes so score blocks stay within this budget.  This
             parameter is appended after the historical positional arguments
             to preserve their calling convention.
+        prototype_refinement : bool, default=False
+            Whether to apply one-pass balanced observation-median refinement
+            after KMeans or MiniBatchKMeans fitting. ``True`` selects the
+            internal ``"balanced_median"`` method; ``False`` leaves fitted
+            prototypes unchanged. Multi-label targets are not supported when
+            refinement is enabled.
         """
         self.rho = rho
         self.r_hat = r_hat
@@ -502,7 +513,9 @@ def __init__(
         self.multilabel_pair_mode = multilabel_pair_mode
         self.top_m = top_m
         self.exclude_classes = exclude_classes
+        self.prototype_refinement = prototype_refinement
         self._validate_multilabel_params()
+        self._validate_prototype_refinement()
 
         # indices / bookkeeping
         self.sparse_adj = defaultdict(int)
@@ -524,6 +537,9 @@ def __init__(
         self._positive_rows_by_label_index_ = {}
         self._score_classes = ()
         self.index = 1.0
+        self.prototype_refinement_ = empty_refinement_summary(
+            refinement_method(self.prototype_refinement)
+        )
 
         self._model: _BaseManyToOneClusteringModel = self._build_model()
 
@@ -549,6 +565,24 @@ def _validate_multilabel_params(self) -> None:
             "offline_memory_budget_mb",
         )
 
+    def _validate_prototype_refinement(self) -> None:
+        """Validate the public boolean refinement switch."""
+        # ``bool`` is deliberately strict here: accepting strings or integer
+        # sentinels would make sklearn cloning and parameter introspection
+        # ambiguous (and ``bool`` is a subclass of ``int`` in Python).
+        refinement_method(self.prototype_refinement)
+        if (
+            self.prototype_refinement
+            and (
+                not isinstance(self.model_type, str)
+                or self.model_type not in {"KMeans", "MiniBatchKMeans"}
+            )
+        ):
+            raise ValueError(
+                "prototype_refinement=True is supported only for "
+                "model_type='KMeans' or 'MiniBatchKMeans'."
+            )
+
     def _build_model(self) -> _BaseManyToOneClusteringModel:
         """Construct the backend adapter from the current estimator parameters."""
         if self.model_type in ["Fuzzy", "Hypersphere"]:
@@ -558,9 +592,21 @@ def _build_model(self) -> _BaseManyToOneClusteringModel:
                 r_hat=self.r_hat,
             )
         if self.model_type == "KMeans":
-            return _KMeansManyToOne(k=self.kmeans_k, kmeans_kwargs=self.kmeans_kwargs)
+            return _KMeansManyToOne(
+                k=self.kmeans_k,
+                kmeans_kwargs=self.kmeans_kwargs,
+                prototype_refinement=self.prototype_refinement,
+                refinement_memory_budget_mb=self.offline_memory_budget_mb,
+                refinement_row_cap=self.offline_chunk_size,
+            )
         if self.model_type == "MiniBatchKMeans":
-            return _MiniBatchKMeansManyToOne(k=self.kmeans_k, kmeans_kwargs=self.kmeans_kwargs)
+            return _MiniBatchKMeansManyToOne(
+                k=self.kmeans_k,
+                kmeans_kwargs=self.kmeans_kwargs,
+                prototype_refinement=self.prototype_refinement,
+                refinement_memory_budget_mb=self.offline_memory_budget_mb,
+                refinement_row_cap=self.offline_chunk_size,
+            )
         if self.model_type == "BallCover":
             kwargs = self.ballcover_kwargs or {}
             return _BallCoverManyToOne(
@@ -572,8 +618,25 @@ def _build_model(self) -> _BaseManyToOneClusteringModel:
 
     def set_params(self, **params: Any) -> "OverlapIndex":
         """Update estimator parameters and rebuild the backend adapter."""
+        # Validate the public switch before BaseEstimator mutates attributes so
+        # a rejected value cannot leave this estimator in a half-updated state.
+        if "prototype_refinement" in params:
+            refinement_method(params["prototype_refinement"])
+            candidate_refinement = params["prototype_refinement"]
+        else:
+            candidate_refinement = self.prototype_refinement
+        candidate_model_type = params.get("model_type", self.model_type)
+        if candidate_refinement and (
+            not isinstance(candidate_model_type, str)
+            or candidate_model_type not in {"KMeans", "MiniBatchKMeans"}
+        ):
+            raise ValueError(
+                "prototype_refinement=True is supported only for "
+                "model_type='KMeans' or 'MiniBatchKMeans'."
+            )
         super().set_params(**params)
         self._validate_multilabel_params()
+        self._validate_prototype_refinement()
         self._model = self._build_model()
         self._reset_indices()
         return self
@@ -711,6 +774,9 @@ def _reset_indices(self) -> None:
         self._positive_rows_by_label_index_ = {}
         self._score_classes = ()
         self.index = 1.0
+        self.prototype_refinement_ = empty_refinement_summary(
+            refinement_method(self.prototype_refinement)
+        )
         if hasattr(self, "n_features_in_"):
             del self.n_features_in_
 
@@ -1067,6 +1133,13 @@ def score_fixed(self, X: np.ndarray, Y: Any) -> float:
             raise ValueError("This OverlapIndex instance is not fit yet.")
 
         X_eval, Y_sets = self._validate_input_data(X, Y)
+        if self.prototype_refinement and any(
+            len(labels) > 1 for labels in Y_sets
+        ):
+            raise ValueError(
+                "prototype_refinement=True does not support "
+                "multi-label targets."
+            )
         self._check_feature_count(X_eval)
         if X_eval.shape[0] == 0:
             self._warn_empty_input()
@@ -1088,7 +1161,11 @@ def score_fixed(self, X: np.ndarray, Y: Any) -> float:
             )
 
         feature_count = int(self.n_features_in_)
+        refinement_summary = self.prototype_refinement_
         self._reset_indices()
+        # ``score_fixed`` recomputes overlap diagnostics but must not discard
+        # the fit-time refinement decisions that describe the held prototypes.
+        self.prototype_refinement_ = refinement_summary
         self.n_features_in_ = feature_count
         self.rev_map = defaultdict(
             set,
@@ -2025,6 +2102,13 @@ def fit_offline(self, X: np.ndarray, Y: Any, reset_state: bool = True) -> float:
             )
 
         X, Y_sets = self._validate_input_data(X, Y)
+        if self.prototype_refinement and any(
+            len(labels) > 1 for labels in Y_sets
+        ):
+            raise ValueError(
+                "prototype_refinement=True does not support "
+                "multi-label targets."
+            )
         if reset_state:
             self._reset_indices()
             self._model = self._build_model()
@@ -2065,6 +2149,13 @@ def fit_offline(self, X: np.ndarray, Y: Any, reset_state: bool = True) -> float:
             )
         else:
             self._model.fit_offline(X_fit, Y_fit)
+        backend_summary = getattr(self._model, "prototype_refinement_summary", None)
+        if backend_summary is None:
+            backend_summary = empty_refinement_summary(
+                refinement_method(self.prototype_refinement),
+                prototype_count=int(self._model.n_clusters_total),
+            )
+        self.prototype_refinement_ = dict(backend_summary)
         self.n_features_in_ = int(X.shape[1])
         self.rev_map = defaultdict(set, {c: set(s) for c, s in self._model.class_to_clusters.items()})
         self._refresh_under_prototyped_labels()