diff --git a/README.md b/README.md index eb1b647..601d669 100644 --- a/README.md +++ b/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. diff --git a/docs/backends/kmeans.md b/docs/backends/kmeans.md index 2d289d6..6db56a8 100644 --- a/docs/backends/kmeans.md +++ b/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. diff --git a/docs/backends/minibatch_kmeans.md b/docs/backends/minibatch_kmeans.md index 27a4554..814fd7c 100644 --- a/docs/backends/minibatch_kmeans.md +++ b/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 diff --git a/docs/diagnostics.md b/docs/diagnostics.md index 2b746ae..1ada264 100644 --- a/docs/diagnostics.md +++ b/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 diff --git a/overlapindex/OverlapIndex.py b/overlapindex/OverlapIndex.py index 29c2a28..4e48916 100644 --- a/overlapindex/OverlapIndex.py +++ b/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() diff --git a/overlapindex/_prototype_refinement.py b/overlapindex/_prototype_refinement.py new file mode 100644 index 0000000..49a9bbb --- /dev/null +++ b/overlapindex/_prototype_refinement.py @@ -0,0 +1,418 @@ +"""Optional one-pass prototype refinement for centroid backends. + +The public estimator keeps this feature deliberately small and opt-in. The +implementation in this module operates on a fitted centroid adapter, computes +outgoing isolation with the universal tiled scorer, and then replaces selected +parent centers with two deterministic observation representatives. It does +not fit a second sklearn model and never materializes a samples-by-prototypes +distance matrix. +""" + +from __future__ import annotations + +from collections import defaultdict +from typing import Any, Dict, Mapping, Optional, Sequence, Tuple + +import numpy as np +from scipy import sparse + +from overlapindex._universal_scorer import ( + compute_second_best_source_scores, + iter_target_class_blocks, +) + + +REFINEMENT_METHOD = "balanced_median" +NO_REFINEMENT_METHOD = "none" + + +def refinement_method(enabled: bool) -> str: + """Return the internal refinement method name for a public bool flag. + + ``prototype_refinement`` is intentionally a strict Python ``bool`` at the + estimator boundary. Keeping this conversion in one place avoids leaking + the public flag into fitted diagnostics, which continue to expose the + descriptive method names used by the refinement implementation. + """ + + if type(enabled) is not bool: + raise ValueError("prototype_refinement must be a boolean (True or False).") + return REFINEMENT_METHOD if enabled else NO_REFINEMENT_METHOD + + +def _plain_label(value: Any) -> Any: + """Convert NumPy scalar labels to ordinary Python values for diagnostics.""" + + return value.item() if isinstance(value, np.generic) else value + + +def empty_refinement_summary( + method: str = NO_REFINEMENT_METHOD, + *, + prototype_count: int = 0, +) -> dict[str, Any]: + """Return the fitted-diagnostics shape used by the estimator. + + ``records`` is a tuple so callers can retain a stable, read-mostly view of + per-parent decisions. Individual records are ordinary dictionaries for + convenient serialization and backwards-compatible key access. + """ + + count = int(prototype_count) + return { + "method": str(method), + "mode": str(method), + "splitter": str(method), + "split_method": str(method), + "prototype_count_before": count, + "prototype_count_after": count, + "eligible_count": 0, + "attempted_count": 0, + "applied_count": 0, + "skipped_count": 0, + "attempted": 0, + "applied": 0, + "skipped": 0, + "eligible_parent_ids": (), + "applied_parent_ids": (), + "skipped_parent_ids": (), + "records": (), + } + + +def _select_rows_dense(X: Any, rows: np.ndarray) -> np.ndarray: + """Select fit rows and densify only the local parent support block.""" + + if sparse.issparse(X): + values = X[rows].toarray() + else: + values = np.asarray(X[rows], dtype=float) + values = np.asarray(values, dtype=float) + if values.ndim != 2: + raise ValueError("prototype refinement requires a two-dimensional feature matrix") + return values + + +def _fit_isolation( + backend: Any, + X: Any, + Y: np.ndarray, + classes: Sequence[Any], + *, + memory_budget_mb: int, + row_cap: Optional[int], +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Return source best ids, source runner-up scores, support and counts. + + ``runner_count`` is the outgoing isolation count used by the research + prototype: a source row contributes only when no competing class's best + prototype strictly beats the source-class second-best score. + """ + + label_to_position = {label: int(position) for position, label in enumerate(classes)} + source_positions = np.asarray( + [label_to_position[label] for label in np.asarray(Y, dtype=object)], + dtype=int, + ) + class_ids = backend.class_center_id_arrays + integer_ids = { + int(position): np.asarray(class_ids.get(label, np.asarray([], dtype=int)), dtype=int).reshape(-1) + for label, position in label_to_position.items() + } + rows = np.arange(int(source_positions.size), dtype=int) + source = compute_second_best_source_scores( + X, + rows, + source_positions, + integer_ids, + backend, + memory_budget_mb=int(memory_budget_mb), + row_cap=row_cap, + ) + best_ids = np.asarray(source.best_prototype_ids, dtype=int) + own_second = np.asarray(source.second_best_scores, dtype=float) + + n_prototypes = int(backend.centers.shape[0]) + support = np.zeros(n_prototypes, dtype=int) + valid_best = best_ids[(best_ids >= 0) & (best_ids < n_prototypes)] + if valid_best.size: + support += np.bincount(valid_best, minlength=n_prototypes)[:n_prototypes] + + competitor_beats = np.zeros(best_ids.size, dtype=bool) + for block in iter_target_class_blocks( + X, + integer_ids, + backend, + memory_budget_mb=int(memory_budget_mb), + row_cap=row_cap, + sample_rows=rows, + ): + block_rows = np.asarray(block.row_indices, dtype=int) + block_sources = source_positions[block_rows] + for column, target_position in enumerate(block.class_ids.tolist()): + target_position = int(target_position) + competitor_beats[block_rows] |= ( + (block_sources != target_position) + & (np.asarray(block.best_scores[:, column]) > own_second[block_rows]) + ) + + runner_count = np.zeros(n_prototypes, dtype=int) + isolated_rows = np.flatnonzero(~competitor_beats) + isolated_ids = best_ids[isolated_rows] + isolated_ids = isolated_ids[(isolated_ids >= 0) & (isolated_ids < n_prototypes)] + if isolated_ids.size: + runner_count += np.bincount(isolated_ids, minlength=n_prototypes)[:n_prototypes] + return best_ids, own_second, support, runner_count + + +def _stable_child_order(children: np.ndarray) -> np.ndarray: + """Return a deterministic lexicographic ordering of two child centers.""" + + return np.asarray( + np.lexsort(tuple(children[:, column] for column in reversed(range(children.shape[1])))), + dtype=int, + ) + + +def _split_parent( + X: Any, + rows: np.ndarray, + parent_center: np.ndarray, +) -> Tuple[Optional[np.ndarray], Optional[Tuple[int, int]], Optional[Tuple[int, int]], str]: + """Build balanced observation-median children for one parent. + + Returns ``(children, selected_rows, child_supports, reason)``. A non-empty + reason denotes a skipped parent; applied splits return an empty reason. + """ + + if rows.size < 2: + return None, None, None, "support<2" + points = _select_rows_dense(X, rows) + if not np.isfinite(points).all(): + return None, None, None, "invalid_points" + parent_center = np.asarray(parent_center, dtype=float).reshape(-1) + if parent_center.size != points.shape[1] or not np.isfinite(parent_center).all(): + return None, None, None, "invalid_parent" + + first_distances = np.sum((points - parent_center) ** 2, axis=1) + if not np.isfinite(first_distances).all(): + return None, None, None, "invalid_points" + first = int(np.argmax(first_distances)) + second_distances = np.sum((points - points[first]) ** 2, axis=1) + if not np.isfinite(second_distances).all(): + return None, None, None, "invalid_points" + second = int(np.argmax(second_distances)) + if second == first or np.array_equal(points[first], points[second]): + return None, None, None, "duplicate_points" + + direction = np.asarray(points[second] - points[first], dtype=float) + direction_norm = float(np.linalg.norm(direction)) + if not np.isfinite(direction_norm) or direction_norm <= 0.0: + return None, None, None, "invalid_axis" + projections = np.asarray(points @ (direction / direction_norm), dtype=float) + if not np.isfinite(projections).all(): + return None, None, None, "invalid_projection" + + # ``rows`` is in original fit-row coordinates and therefore makes exact + # projection ties deterministic regardless of sparse/dense slicing order. + ordered = np.lexsort((rows, projections)) + split_at = int(points.shape[0] // 2) + if split_at <= 0 or split_at >= points.shape[0]: + return None, None, None, "empty_child" + halves = (ordered[:split_at], ordered[split_at:]) + representatives = [] + representative_rows = [] + child_supports = [] + for half in halves: + half_points = points[half] + half_rows = rows[half] + coordinate_median = np.median(half_points, axis=0) + squared_distances = np.sum((half_points - coordinate_median) ** 2, axis=1) + if not np.isfinite(coordinate_median).all() or not np.isfinite(squared_distances).all(): + return None, None, None, "invalid_child" + nearest_order = np.lexsort((half_rows, squared_distances)) + nearest = int(nearest_order[0]) + representatives.append(np.asarray(half_points[nearest], dtype=float)) + representative_rows.append(int(half_rows[nearest])) + child_supports.append(int(half.size)) + + children = np.asarray(representatives, dtype=float) + if ( + children.shape != (2, points.shape[1]) + or not np.isfinite(children).all() + or np.array_equal(children[0], children[1]) + ): + return None, None, None, "duplicate_representatives" + + order = _stable_child_order(children) + children = children[order] + representative_rows = [representative_rows[int(index)] for index in order.tolist()] + child_supports = [child_supports[int(index)] for index in order.tolist()] + return ( + children, + (int(representative_rows[0]), int(representative_rows[1])), + (int(child_supports[0]), int(child_supports[1])), + "", + ) + + +def apply_balanced_median_refinement( + backend: Any, + X: Any, + Y: Any, + *, + memory_budget_mb: int = 256, + row_cap: Optional[int] = None, +) -> dict[str, Any]: + """Apply one frozen balanced observation-median pass to a backend. + + The adapter has already been fitted when this function is called. Parent + eligibility is measured once against the original center set; appended + children are never reconsidered during the same pass. + """ + + centers_before = np.asarray(backend.centers) + prototype_count_before = int(centers_before.shape[0]) + summary = empty_refinement_summary( + REFINEMENT_METHOD, + prototype_count=prototype_count_before, + ) + if prototype_count_before == 0: + return summary + + Y_array = np.asarray(Y, dtype=object).reshape(-1) + if Y_array.size == 0: + return summary + classes = list(dict.fromkeys(Y_array.tolist())) + if len(classes) < 2: + return summary + + best_ids, _own_second, support, runner_count = _fit_isolation( + backend, + X, + Y_array, + classes, + memory_budget_mb=int(memory_budget_mb), + row_cap=row_cap, + ) + owner_values = np.asarray(backend.cluster_to_class, dtype=object) + owner_for = { + int(pid): owner_values[int(pid)] + for pid in range(min(prototype_count_before, owner_values.size)) + } + eligible = tuple( + int(pid) + for pid in range(prototype_count_before) + if int(support[pid]) >= 2 and int(runner_count[pid]) == 0 + ) + + records = [] + applied_parent_ids = [] + skipped_parent_ids = [] + replacements: Dict[int, np.ndarray] = {} + appended: list[tuple[int, Any, np.ndarray]] = [] + next_id = prototype_count_before + + original_centers = np.asarray(centers_before, dtype=float) + for parent in eligible: + rows = np.flatnonzero(best_ids == parent).astype(int, copy=False) + label = _plain_label(owner_for.get(parent)) + children, selected_rows, child_supports, reason = _split_parent( + X, + rows, + original_centers[parent], + ) + base_record: dict[str, Any] = { + "parent": int(parent), + "parent_id": int(parent), + "original_id": int(parent), + "class": label, + "class_label": label, + "support": int(rows.size), + "support_before": int(rows.size), + "status": "skipped", + "reason": str(reason), + "children": (), + "child_ids": (), + "child_supports": (), + "selected_observation_indices": (), + } + if children is None: + skipped_parent_ids.append(int(parent)) + records.append(base_record) + continue + + child_id = int(next_id) + next_id += 1 + replacements[int(parent)] = np.asarray(children[0], dtype=backend._dtype) + appended.append((child_id, label, np.asarray(children[1], dtype=backend._dtype))) + applied_parent_ids.append(int(parent)) + base_record.update( + { + "status": "applied", + "reason": "", + "children": (int(parent), child_id), + "child_ids": (int(parent), child_id), + "child_supports": tuple(int(value) for value in child_supports or ()), + "selected_observation_indices": tuple(int(value) for value in selected_rows or ()), + "selected_sample_indices": tuple(int(value) for value in selected_rows or ()), + "new_id": child_id, + } + ) + records.append(base_record) + + if appended: + final_centers = np.array(original_centers, copy=True) + for parent, replacement in replacements.items(): + final_centers[parent] = replacement + final_centers = np.vstack( + [final_centers] + [child for _child_id, _label, child in appended] + ).astype(backend._dtype, copy=False) + backend._centers = final_centers + backend._center_norms = np.einsum("ij,ij->i", final_centers, final_centers) + + class_lists = { + label: [int(value) for value in np.asarray(ids, dtype=int).reshape(-1).tolist()] + for label, ids in backend._class_center_ids.items() + } + for child_id, label, _child in appended: + class_lists.setdefault(label, []).append(int(child_id)) + backend._class_center_ids = class_lists + backend._class_center_id_arrays = { + label: np.asarray(ids, dtype=int) for label, ids in class_lists.items() + } + backend._class_to_clusters = defaultdict( + set, + {label: set(ids) for label, ids in class_lists.items()}, + ) + original_owners = np.asarray(backend._cluster_to_class, dtype=object).reshape(-1) + appended_owners = np.asarray([label for _child_id, label, _child in appended], dtype=object) + backend._cluster_to_class = np.concatenate((original_owners, appended_owners)) + + prototype_count_after = int(backend._centers.shape[0]) + summary.update( + { + "prototype_count_after": prototype_count_after, + "eligible_count": len(eligible), + "attempted_count": len(eligible), + "applied_count": len(applied_parent_ids), + "skipped_count": len(skipped_parent_ids), + "attempted": len(eligible), + "applied": len(applied_parent_ids), + "skipped": len(skipped_parent_ids), + "eligible_parent_ids": tuple(eligible), + "applied_parent_ids": tuple(applied_parent_ids), + "skipped_parent_ids": tuple(skipped_parent_ids), + "records": tuple(dict(record) for record in records), + } + ) + return summary + + +__all__ = [ + "REFINEMENT_METHOD", + "NO_REFINEMENT_METHOD", + "refinement_method", + "empty_refinement_summary", + "apply_balanced_median_refinement", +] diff --git a/overlapindex/clustering.py b/overlapindex/clustering.py index 65c339a..2335921 100644 --- a/overlapindex/clustering.py +++ b/overlapindex/clustering.py @@ -9,6 +9,11 @@ _validate_class_dictionary_coverage, _validate_positive_integer, ) +from overlapindex._prototype_refinement import ( + apply_balanced_median_refinement, + empty_refinement_summary, + refinement_method, +) from typing import Literal, Optional, Union, Dict, Any, Sequence, Tuple, Type @@ -220,6 +225,9 @@ def __init__( k: Union[int, Dict[Any, int]] = 8, model_kwargs: Optional[dict] = None, dtype: Type[np.floating] = np.float32, + prototype_refinement: bool = False, + refinement_memory_budget_mb: int = 256, + refinement_row_cap: Optional[int] = None, ) -> None: """ Initialize shared centroid-backend state. @@ -232,6 +240,14 @@ def __init__( Keyword arguments forwarded to the concrete clustering estimator. dtype : numpy floating dtype, default=np.float32 Floating-point dtype used to store centroid arrays. + prototype_refinement : bool, default=False + Whether to apply one-pass observation-median refinement after the + initial centroid fit. ``True`` selects the internal + ``"balanced_median"`` method. + refinement_memory_budget_mb : int, default=256 + Scratch-memory budget used by the tiled fit-isolation pass. + refinement_row_cap : int or None, optional + Optional row cap for the tiled fit-isolation pass. """ self._k = k if isinstance(k, dict): @@ -241,6 +257,11 @@ def __init__( _validate_positive_integer(k, "k") self._model_kwargs = model_kwargs or {} self._dtype = dtype + # Keep a descriptive resolved method name in fitted diagnostics while + # accepting the estimator's strict boolean switch at its boundary. + self._prototype_refinement = refinement_method(prototype_refinement) + self._refinement_memory_budget_mb = int(refinement_memory_budget_mb) + self._refinement_row_cap = refinement_row_cap self._models: Dict[Any, Any] = {} self._centers: Optional[np.ndarray] = None @@ -249,6 +270,9 @@ def __init__( self._class_center_id_arrays: Dict[Any, np.ndarray] = {} self._class_to_clusters: Dict[Any, set] = defaultdict(set) self._cluster_to_class: Optional[np.ndarray] = None + self._prototype_refinement_summary: dict[str, Any] = empty_refinement_summary( + self._prototype_refinement + ) def _make_model(self, n_clusters: int) -> Any: """Create a concrete centroid estimator with the requested cluster count.""" @@ -313,6 +337,22 @@ def fit_offline(self, X: np.ndarray, Y: np.ndarray) -> None: ) self._center_norms = np.einsum("ij,ij->i", self._centers, self._centers) self._cluster_to_class = np.asarray(cluster_classes, dtype=object) + self._prototype_refinement_summary = empty_refinement_summary( + self._prototype_refinement, + prototype_count=int(self._centers.shape[0]), + ) + if ( + self._prototype_refinement == "balanced_median" + and len(classes) >= 2 + and self._centers.shape[0] > 0 + ): + self._prototype_refinement_summary = apply_balanced_median_refinement( + self, + X, + Y, + memory_budget_mb=self._refinement_memory_budget_mb, + row_cap=self._refinement_row_cap, + ) def partial_fit(self, X: np.ndarray, Y: np.ndarray, **kwargs: Any) -> None: """Raise because centroid backends in this adapter are offline-only.""" @@ -528,6 +568,11 @@ def n_clusters_total(self) -> int: """Return the number of global centroids.""" return 0 if self._centers is None else int(self._centers.shape[0]) + @property + def prototype_refinement_summary(self) -> dict[str, Any]: + """Return fitted diagnostics for the optional refinement pass.""" + return self._prototype_refinement_summary + class _KMeansManyToOne(_BaseCentroidManyToOne): """ @@ -539,11 +584,21 @@ def __init__( self, k: Union[int, Dict[Any, int]] = 8, kmeans_kwargs: Optional[dict] = None, + prototype_refinement: bool = False, + refinement_memory_budget_mb: int = 256, + refinement_row_cap: Optional[int] = None, ) -> None: """Initialize a per-class scikit-learn KMeans backend.""" if KMeans is None: raise ImportError("scikit-learn is required for model_type='KMeans'.") - super().__init__(k=k, model_kwargs=kmeans_kwargs, dtype=np.float32) + super().__init__( + k=k, + model_kwargs=kmeans_kwargs, + dtype=np.float32, + prototype_refinement=prototype_refinement, + refinement_memory_budget_mb=refinement_memory_budget_mb, + refinement_row_cap=refinement_row_cap, + ) def _make_model(self, n_clusters: int) -> KMeans: """Create a scikit-learn KMeans estimator.""" @@ -563,11 +618,21 @@ def __init__( self, k: Union[int, Dict[Any, int]] = 8, kmeans_kwargs: Optional[dict] = None, + prototype_refinement: bool = False, + refinement_memory_budget_mb: int = 256, + refinement_row_cap: Optional[int] = None, ) -> None: """Initialize a per-class scikit-learn MiniBatchKMeans backend.""" if MiniBatchKMeans is None: raise ImportError("scikit-learn is required for model_type='MiniBatchKMeans'.") - super().__init__(k=k, model_kwargs=kmeans_kwargs, dtype=np.float32) + super().__init__( + k=k, + model_kwargs=kmeans_kwargs, + dtype=np.float32, + prototype_refinement=prototype_refinement, + refinement_memory_budget_mb=refinement_memory_budget_mb, + refinement_row_cap=refinement_row_cap, + ) def _make_model(self, n_clusters: int) -> MiniBatchKMeans: diff --git a/tests/test_prototype_refinement.py b/tests/test_prototype_refinement.py new file mode 100644 index 0000000..d7e13e2 --- /dev/null +++ b/tests/test_prototype_refinement.py @@ -0,0 +1,481 @@ +"""Public contracts for the optional balanced-median prototype refinement.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import numpy as np +import pytest +from scipy import sparse +from sklearn.base import clone + +from overlapindex import OverlapIndex + + +BACKENDS = ("KMeans", "MiniBatchKMeans") + + +def _kmeans_kwargs(model_type: str) -> dict[str, object]: + kwargs: dict[str, object] = {"random_state": 0, "n_init": 10} + if model_type == "MiniBatchKMeans": + kwargs.update({"batch_size": 8, "max_iter": 100}) + return kwargs + + +def _model( + model_type: str, + *, + refinement: bool = False, + k: int = 1, + **kwargs: object, +) -> OverlapIndex: + params: dict[str, object] = { + "model_type": model_type, + "kmeans_k": k, + "kmeans_kwargs": _kmeans_kwargs(model_type), + "prototype_refinement": refinement, + } + params.update(kwargs) + return OverlapIndex(**params) + + +def _balanced_data() -> tuple[np.ndarray, np.ndarray]: + """One-dimensional classes whose median halves have known observations.""" + + X = np.asarray( + [[0.0], [1.0], [2.0], [3.0], [10.0], [11.0], [12.0], [13.0]], + dtype=float, + ) + y = np.asarray([0, 0, 0, 0, 1, 1, 1, 1]) + return X, y + + +def _multiclass_data() -> tuple[np.ndarray, np.ndarray]: + X = np.asarray( + [ + [0.0], + [1.0], + [2.0], + [3.0], + [10.0], + [11.0], + [12.0], + [13.0], + [20.0], + [21.0], + [22.0], + [23.0], + ], + dtype=float, + ) + # Non-contiguous labels catch accidental use of labels as global IDs. + y = np.asarray([10] * 4 + [20] * 4 + [30] * 4) + return X, y + + +def _trap_data() -> tuple[np.ndarray, np.ndarray]: + """Two prototypes per class with every parent eligible for a split.""" + + X = np.asarray( + [ + [0.0], + [0.1], + [10.0], + [9.9], + [3.0], + [3.1], + [7.0], + [6.9], + ], + dtype=float, + ) + y = np.asarray([0] * 4 + [1] * 4) + return X, y + + +def _far_prototypes_data() -> tuple[np.ndarray, np.ndarray]: + """Own runner-up scores dominate distant competitors, so no parent splits.""" + + X = np.asarray( + [ + [0.0], + [0.1], + [10.0], + [10.1], + [100.0], + [100.1], + [110.0], + [110.1], + ], + dtype=float, + ) + y = np.asarray([0] * 4 + [1] * 4) + return X, y + + +def _duplicate_data() -> tuple[np.ndarray, np.ndarray]: + X = np.asarray([[0.0], [0.0], [0.0], [0.0], [10.0], [10.0], [10.0], [10.0]]) + y = np.asarray([0] * 4 + [1] * 4) + return X, y + + +def _summary(model: OverlapIndex) -> Mapping[str, object]: + value = getattr(model, "prototype_refinement_", None) + assert isinstance(value, Mapping), "fitted prototype_refinement_ summary is required" + return value + + +def _records(model: OverlapIndex) -> tuple[Mapping[str, object], ...]: + records = _summary(model).get("records") + assert records is not None + return tuple(records) # type: ignore[arg-type] + + +def _backend_centers(model: OverlapIndex) -> np.ndarray: + return np.asarray(model._model.centers) + + +def _backend_ids(model: OverlapIndex) -> dict[object, np.ndarray]: + return { + label: np.asarray(values, dtype=int) + for label, values in model._model.class_center_id_arrays.items() + } + + +def _assert_mapping_close(left: Mapping[object, object], right: Mapping[object, object]) -> None: + assert set(left) == set(right) + for key in left: + lhs, rhs = left[key], right[key] + if isinstance(lhs, (float, np.floating)) or isinstance(rhs, (float, np.floating)): + assert float(lhs) == pytest.approx(float(rhs), abs=1e-7) + else: + assert lhs == rhs + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_default_off_and_explicit_false_are_backward_parity(model_type: str) -> None: + X, y = _balanced_data() + default = OverlapIndex( + model_type=model_type, + kmeans_k=1, + kmeans_kwargs=_kmeans_kwargs(model_type), + ).fit(X, y) + explicit = _model(model_type, refinement=False).fit(X, y) + + assert default.prototype_refinement is False + assert explicit.prototype_refinement is False + assert default.get_params() == explicit.get_params() + assert default.index == pytest.approx(explicit.index, abs=0.0) + assert default.weighted_index == pytest.approx(explicit.weighted_index, abs=0.0) + _assert_mapping_close(default.singleton_index, explicit.singleton_index) + assert dict(default.cluster_cardinality) == dict(explicit.cluster_cardinality) + assert {label: set(ids) for label, ids in default.rev_map.items()} == { + label: set(ids) for label, ids in explicit.rev_map.items() + } + np.testing.assert_array_equal(_backend_centers(default), _backend_centers(explicit)) + np.testing.assert_array_equal(default.predict(X), explicit.predict(X)) + + summary = _summary(default) + assert summary["method"] == "none" + assert summary["eligible_count"] == 0 + assert summary["attempted_count"] == 0 + assert summary["applied_count"] == 0 + assert summary["skipped_count"] == 0 + assert summary["prototype_count_before"] == summary["prototype_count_after"] + assert tuple(summary["records"]) == () + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_true_balanced_median_is_deterministic_one_pass_and_observation_based( + model_type: str, +) -> None: + X, y = _balanced_data() + first = _model(model_type, refinement=True).fit(X, y) + second = _model(model_type, refinement=True).fit(X, y) + + assert first.prototype_refinement is True + assert second.prototype_refinement is True + np.testing.assert_array_equal(_backend_centers(first), _backend_centers(second)) + assert _backend_ids(first).keys() == _backend_ids(second).keys() + for label in _backend_ids(first): + np.testing.assert_array_equal(_backend_ids(first)[label], _backend_ids(second)[label]) + assert dict(_summary(first)) == dict(_summary(second)) + + summary = _summary(first) + assert summary["method"] == "balanced_median" + assert summary["prototype_count_before"] == 2 + assert summary["prototype_count_after"] == 4 + assert summary["eligible_count"] == 2 + assert summary["attempted_count"] == 2 + assert summary["applied_count"] == 2 + assert summary["skipped_count"] == 0 + + centers = _backend_centers(first) + records = _records(first) + assert len(records) == 2 + assert { + int(record["parent_id"]): tuple(record["selected_observation_indices"]) + for record in records + } == {0: (0, 2), 1: (4, 6)} + for record in records: + assert record["status"] == "applied" + assert int(record["support"]) >= 2 + children = np.asarray(record["child_ids"], dtype=int) + assert children.shape == (2,) + supports = np.asarray(record["child_supports"], dtype=int) + assert supports.shape == (2,) + assert np.all(supports > 0) + assert int(supports.sum()) == int(record["support"]) + selected = np.asarray(record["selected_observation_indices"], dtype=int) + assert selected.shape == (2,) + # Every child center is an actual fit observation, never a synthetic + # coordinate median; the integer fixture also makes this exact. + np.testing.assert_array_equal(centers[children], X[selected]) + + # One pass: children are not recursively eligible in the same fit. + assert summary["prototype_count_after"] == summary["prototype_count_before"] + summary["applied_count"] + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_eligibility_uses_support_threshold_and_outgoing_runner_up(model_type: str) -> None: + # Support exactly two is eligible with one prototype per class. + X_two = np.asarray([[0.0], [1.0], [10.0], [11.0]]) + y_two = np.asarray([0, 0, 1, 1]) + two = _model(model_type, refinement=True).fit(X_two, y_two) + assert _summary(two)["eligible_count"] == 2 + assert _summary(two)["applied_count"] == 2 + assert _summary(two)["prototype_count_after"] == 4 + + # A one-row parent is below the fixed support >=2 eligibility threshold; + # the other class still refines normally. + X_one = np.asarray([[0.0], [10.0], [11.0], [12.0], [13.0]]) + y_one = np.asarray([0, 1, 1, 1, 1]) + one = _model(model_type, refinement=True).fit(X_one, y_one) + assert _summary(one)["eligible_count"] == 1 + assert _summary(one)["applied_count"] == 1 + assert _backend_ids(one)[0].size == 1 + assert _backend_ids(one)[1].size == 2 + + # With well-separated two-prototype classes, own runner-up evidence is + # present for every parent; support alone must not make them eligible. + X_far, y_far = _far_prototypes_data() + far = _model(model_type, refinement=True, k=2).fit(X_far, y_far) + assert _summary(far)["eligible_count"] == 0 + assert _summary(far)["applied_count"] == 0 + assert _summary(far)["prototype_count_before"] == 4 + assert _summary(far)["prototype_count_after"] == 4 + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_multiple_parents_preserve_class_ownership_and_contiguous_global_ids( + model_type: str, +) -> None: + X, y = _trap_data() + model = _model(model_type, refinement=True, k=2).fit(X, y) + summary = _summary(model) + ids = _backend_ids(model) + + assert summary["prototype_count_before"] == 4 + assert summary["prototype_count_after"] == 8 + assert summary["applied_count"] == 4 + all_ids = sorted(int(pid) for values in ids.values() for pid in values) + assert all_ids == list(range(8)) + # Original IDs stay in place; newly created children are appended to the + # global ID range rather than interleaved into each class's old block. + prototype_count_before = int(summary["prototype_count_before"]) + records = _records(model) + for record in records: + assert record["child_ids"][0] == record["parent_id"] + assert int(record["child_ids"][1]) >= prototype_count_before + appended_ids = sorted( + int(record["child_ids"][1]) + for record in records + if record["status"] == "applied" + ) + assert appended_ids == list(range(prototype_count_before, int(summary["prototype_count_after"]))) + cluster_to_class = np.asarray(model._model.cluster_to_class, dtype=object) + assert cluster_to_class.shape == (8,) + for label, values in ids.items(): + assert set(cluster_to_class[values].tolist()) == {label} + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_duplicate_observation_representatives_are_skipped_without_id_gaps( + model_type: str, +) -> None: + X, y = _duplicate_data() + model = _model(model_type, refinement=True).fit(X, y) + summary = _summary(model) + + assert summary["eligible_count"] == 2 + assert summary["attempted_count"] == 2 + assert summary["applied_count"] == 0 + assert summary["skipped_count"] == 2 + assert summary["prototype_count_before"] == summary["prototype_count_after"] == 2 + assert all(record["status"] == "skipped" for record in _records(model)) + assert all( + str(record["reason"]) + in {"duplicate_points", "duplicate_representatives", "no_progress", "invalid_child"} + for record in _records(model) + ) + ids = _backend_ids(model) + assert sorted(int(pid) for values in ids.values() for pid in values) == [0, 1] + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_multiclass_noncontiguous_labels_keep_class_owned_children(model_type: str) -> None: + X, y = _multiclass_data() + model = _model(model_type, refinement=True).fit(X, y) + ids = _backend_ids(model) + assert set(ids) == {10, 20, 30} + assert sorted(int(pid) for values in ids.values() for pid in values) == list(range(6)) + for label, values in ids.items(): + assert values.size == 2 + assert set(np.asarray(model._model.cluster_to_class, dtype=object)[values].tolist()) == {label} + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_score_fixed_is_immutable_and_does_not_run_a_second_refinement(model_type: str) -> None: + X, y = _balanced_data() + model = _model(model_type, refinement=True).fit(X, y) + centers_before = _backend_centers(model).copy() + ids_before = {label: values.copy() for label, values in _backend_ids(model).items()} + summary_before = dict(_summary(model)) + eval_X = np.asarray([[0.5], [2.5], [10.5], [12.5]]) + eval_y = np.asarray([0, 0, 1, 1]) + + score = model.score_fixed(eval_X, eval_y) + + assert np.isfinite(score) + np.testing.assert_array_equal(_backend_centers(model), centers_before) + for label, values in ids_before.items(): + np.testing.assert_array_equal(_backend_ids(model)[label], values) + assert dict(_summary(model)) == summary_before + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_repeated_fit_and_score_rebuild_one_refinement_pass(model_type: str) -> None: + X, y = _balanced_data() + model = _model(model_type, refinement=True) + model.fit(X, y) + first_centers = _backend_centers(model).copy() + first_summary = dict(_summary(model)) + + model.fit(X, y) + np.testing.assert_array_equal(_backend_centers(model), first_centers) + assert dict(_summary(model)) == first_summary + model.score(X, y) + np.testing.assert_array_equal(_backend_centers(model), first_centers) + assert dict(_summary(model)) == first_summary + + with pytest.raises(ValueError, match="reset_state=False is supported only for ARTMAP"): + model.fit_offline(X, y, reset_state=False) + np.testing.assert_array_equal(_backend_centers(model), first_centers) + assert dict(_summary(model)) == first_summary + + +def test_get_set_params_and_clone_expose_only_the_public_refinement_switch() -> None: + model = _model("KMeans", refinement=True) + params = model.get_params() + assert params["prototype_refinement"] is True + assert "splitter" not in params + assert "init" not in params + assert "min_support" not in params + + copied = clone(model) + assert copied is not model + assert copied.prototype_refinement is True + assert copied.get_params() == params + copied.fit(*_balanced_data()) + assert _summary(copied)["method"] == "balanced_median" + + model.set_params(prototype_refinement=False) + assert model.prototype_refinement is False + assert model.get_params()["prototype_refinement"] is False + + +def test_invalid_refinement_values_and_unsupported_backends_fail_clearly() -> None: + for invalid in ( + "none", + "balanced_median", + "unknown", + None, + [], + 0, + 1, + np.bool_(True), + ): + with pytest.raises(ValueError, match="prototype_refinement"): + OverlapIndex(prototype_refinement=invalid) + + # The switch is accepted on every backend while disabled; only enabling + # refinement on an unsupported backend should fail. + disabled = OverlapIndex(model_type="BallCover", prototype_refinement=False).fit( + *_balanced_data() + ) + assert disabled.prototype_refinement is False + + with pytest.raises((ValueError, NotImplementedError), match="KMeans|MiniBatchKMeans|prototype_refinement"): + bad = OverlapIndex(model_type="BallCover", prototype_refinement=True) + bad.fit(*_balanced_data()) + + model = _model("KMeans") + with pytest.raises(ValueError, match="prototype_refinement"): + model.set_params(prototype_refinement="unknown") + assert model.prototype_refinement is False + + with pytest.raises(ValueError, match="KMeans|MiniBatchKMeans|prototype_refinement"): + model.set_params(model_type="BallCover", prototype_refinement=True) + assert model.model_type == "KMeans" + assert model.prototype_refinement is False + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_true_balanced_multilabel_rejected_but_false_preserves_existing_behavior(model_type: str) -> None: + X = np.asarray([[0.0], [0.1], [1.0], [1.1], [2.0], [2.1]]) + labels = [{"A", "B"}, {"A"}, {"B"}, {"B", "C"}, {"C"}, {"A", "C"}] + + default = OverlapIndex( + model_type=model_type, + kmeans_k=1, + kmeans_kwargs=_kmeans_kwargs(model_type), + ).fit(X, labels) + explicit = _model(model_type, refinement=False).fit(X, labels) + assert default.index == pytest.approx(explicit.index, abs=0.0) + _assert_mapping_close(default.singleton_index, explicit.singleton_index) + assert dict(default.cluster_cardinality) == dict(explicit.cluster_cardinality) + np.testing.assert_array_equal(_backend_centers(default), _backend_centers(explicit)) + + with pytest.raises((ValueError, NotImplementedError), match="multi-label|multilabel|single-label"): + _model(model_type, refinement=True).fit(X, labels) + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_balanced_score_fixed_rejects_multilabel_evaluation(model_type: str) -> None: + X, y = _balanced_data() + model = _model(model_type, refinement=True).fit(X, y) + eval_X = np.asarray([[0.5], [2.5], [10.5], [12.5]]) + eval_labels = [{0}, {0, 1}, {1}, {1}] + + with pytest.raises((ValueError, NotImplementedError), match="multi-label|multilabel|single-label"): + model.score_fixed(eval_X, eval_labels) + + +@pytest.mark.parametrize("model_type", BACKENDS) +def test_true_balanced_median_sparse_fit_matches_dense_contract(model_type: str) -> None: + X, y = _balanced_data() + dense = _model(model_type, refinement=True).fit(X, y) + sparse_model = _model(model_type, refinement=True).fit( + sparse.csr_matrix(X), y + ) + + assert sparse_model.index == pytest.approx(dense.index, abs=1e-6) + assert sparse_model.weighted_index == pytest.approx(dense.weighted_index, abs=1e-6) + np.testing.assert_allclose(_backend_centers(sparse_model), _backend_centers(dense), atol=1e-6, rtol=0.0) + for label in _backend_ids(dense): + np.testing.assert_array_equal(_backend_ids(sparse_model)[label], _backend_ids(dense)[label]) + assert dict(_summary(sparse_model)) == dict(_summary(dense)) + np.testing.assert_array_equal( + sparse_model.predict(sparse.csr_matrix(X)), + dense.predict(X), + )