from __future__ import annotations from dataclasses import dataclass from typing import Any import numpy as np import pandas as pd from sklearn.cluster import KMeans from sklearn.decomposition import TruncatedSVD from sklearn.metrics import silhouette_score @dataclass(frozen=True) class MissingnessGroupingResult: selection_groups: pd.DataFrame calibration_groups: pd.DataFrame test_groups: pd.DataFrame group_labels: dict[int, str] metadata: dict[str, Any] def build_missingness_groups( *, selection_frame: pd.DataFrame, calibration_frame: pd.DataFrame, test_frame: pd.DataFrame, strategy: str, candidate_missing_variables: list[str] | None = None, mask_cluster_k_grid: list[int] | None = None, min_group_fraction: float = 0.10, min_selection_group_rows: int = 1, random_state: int = 0, **_: Any, ) -> MissingnessGroupingResult: if strategy == "coverage_gap_variable": return _build_coverage_gap_variable_groups( selection_frame=selection_frame, calibration_frame=calibration_frame, test_frame=test_frame, candidate_missing_variables=candidate_missing_variables, min_group_fraction=min_group_fraction, min_selection_group_rows=min_selection_group_rows, ) if strategy == "mask_cluster": return _build_mask_cluster_groups( selection_frame=selection_frame, calibration_frame=calibration_frame, test_frame=test_frame, mask_cluster_k_grid=mask_cluster_k_grid or [2, 3, 4, 5], min_group_fraction=min_group_fraction, min_selection_group_rows=min_selection_group_rows, random_state=random_state, ) raise ValueError(f"unsupported grouping strategy: {strategy}") def _missing_rate_columns(frame: pd.DataFrame) -> list[str]: return [ column for column in frame.columns if column.endswith("_missing_rate") and column != "global_missing_rate" ] def _candidate_missing_rate_columns( frame: pd.DataFrame, candidate_missing_variables: list[str] | None, ) -> list[str]: available = set(_missing_rate_columns(frame)) if candidate_missing_variables is None: return sorted(available) columns = [] for variable in candidate_missing_variables: column = f"{variable}_missing_rate" if column in available: columns.append(column) return columns def _never_observed_indicator(frame: pd.DataFrame, column: str) -> np.ndarray: values = frame[column].to_numpy(dtype=float) return values > 0.0 def _frame_groups_from_ids(group_ids: np.ndarray, group_labels: dict[int, str]) -> pd.DataFrame: return pd.DataFrame( { "group": group_ids.astype(int), "group_label": [group_labels[int(group_id)] for group_id in group_ids], } ) def _build_coverage_gap_variable_groups( *, selection_frame: pd.DataFrame, calibration_frame: pd.DataFrame, test_frame: pd.DataFrame, candidate_missing_variables: list[str] | None, min_group_fraction: float, min_selection_group_rows: int, ) -> MissingnessGroupingResult: candidate_columns = _candidate_missing_rate_columns(selection_frame, candidate_missing_variables) diagnostics: dict[str, dict[str, float | int]] = {} candidate_order = {column: index for index, column in enumerate(candidate_columns)} def collect_diagnostics(*, relaxed: bool) -> dict[str, dict[str, float | int]]: local_diagnostics: dict[str, dict[str, float | int]] = {} for column in candidate_columns: selection_missing = _never_observed_indicator(selection_frame, column) calibration_missing = _never_observed_indicator(calibration_frame, column) selection_fraction = float(selection_missing.mean()) if not relaxed and ( selection_fraction < min_group_fraction or selection_fraction > (1.0 - min_group_fraction) ): continue selection_missing_rows = int(selection_missing.sum()) selection_observed_rows = int((~selection_missing).sum()) calibration_missing_rows = int(calibration_missing.sum()) calibration_observed_rows = int((~calibration_missing).sum()) if not relaxed and min(selection_missing_rows, selection_observed_rows) < min_selection_group_rows: continue calibration_fraction = float(calibration_missing.mean()) selection_gap = abs(selection_fraction - calibration_fraction) imbalance = abs(0.5 - selection_fraction) local_diagnostics[column] = { "selection_missing_fraction": selection_fraction, "calibration_missing_fraction": calibration_fraction, "selection_missing_rows": selection_missing_rows, "selection_observed_rows": selection_observed_rows, "calibration_missing_rows": calibration_missing_rows, "calibration_observed_rows": calibration_observed_rows, "selection_gap": selection_gap, "imbalance": imbalance, "minority_support": min( selection_missing_rows, selection_observed_rows, ), "relaxed_selection": relaxed, } return local_diagnostics diagnostics = collect_diagnostics(relaxed=False) if not diagnostics and candidate_missing_variables is None: diagnostics = collect_diagnostics(relaxed=True) if not diagnostics: raise ValueError("no candidates satisfied minimum support and missing-fraction requirements") for column in candidate_columns: selection_missing = _never_observed_indicator(selection_frame, column) calibration_missing = _never_observed_indicator(calibration_frame, column) selection_fraction = float(selection_missing.mean()) if ( selection_fraction < min_group_fraction or selection_fraction > (1.0 - min_group_fraction) ): continue selection_missing_rows = int(selection_missing.sum()) selection_observed_rows = int((~selection_missing).sum()) calibration_missing_rows = int(calibration_missing.sum()) calibration_observed_rows = int((~calibration_missing).sum()) if min(selection_missing_rows, selection_observed_rows) < min_selection_group_rows: continue calibration_fraction = float(calibration_missing.mean()) selection_gap = abs(selection_fraction - calibration_fraction) imbalance = abs(0.5 - selection_fraction) if column not in diagnostics: continue diagnostics[column].update( { "selection_missing_fraction": selection_fraction, "calibration_missing_fraction": calibration_fraction, "selection_missing_rows": selection_missing_rows, "selection_observed_rows": selection_observed_rows, "calibration_missing_rows": calibration_missing_rows, "calibration_observed_rows": calibration_observed_rows, "selection_gap": selection_gap, "imbalance": imbalance, "minority_support": min( selection_missing_rows, selection_observed_rows, ), } ) selected_column = sorted( diagnostics, key=lambda column: ( -float(diagnostics[column]["selection_gap"]), float(diagnostics[column]["imbalance"]), -int(diagnostics[column]["minority_support"]), candidate_order.get(column, len(candidate_order)), ), )[0] variable = selected_column[: -len("_missing_rate")] group_labels = { 0: f"{variable.lower()}_ever_observed", 1: f"{variable.lower()}_never_observed", } selection_ids = _never_observed_indicator(selection_frame, selected_column).astype(int) calibration_ids = _never_observed_indicator(calibration_frame, selected_column).astype(int) test_ids = _never_observed_indicator(test_frame, selected_column).astype(int) return MissingnessGroupingResult( selection_groups=_frame_groups_from_ids(selection_ids, group_labels), calibration_groups=_frame_groups_from_ids(calibration_ids, group_labels), test_groups=_frame_groups_from_ids(test_ids, group_labels), group_labels=group_labels, metadata={ "strategy": "coverage_gap_variable", "selected_variable": variable, "group_source": "missing_rate_column", "group_count": len(group_labels), "selection_variable_diagnostics": { column[: -len("_missing_rate")]: values for column, values in diagnostics.items() }, }, ) def _mask_feature_matrix(frame: pd.DataFrame, columns: list[str]) -> np.ndarray: return np.column_stack([_never_observed_indicator(frame, column).astype(float) for column in columns]) def _filter_mask_columns( selection_frame: pd.DataFrame, columns: list[str], *, min_group_fraction: float, ) -> list[str]: filtered = [] for column in columns: fraction = float(_never_observed_indicator(selection_frame, column).mean()) if min_group_fraction <= fraction <= (1.0 - min_group_fraction): filtered.append(column) return filtered def _build_mask_cluster_groups( *, selection_frame: pd.DataFrame, calibration_frame: pd.DataFrame, test_frame: pd.DataFrame, mask_cluster_k_grid: list[int], min_group_fraction: float, min_selection_group_rows: int, random_state: int, ) -> MissingnessGroupingResult: columns = _filter_mask_columns( selection_frame, _missing_rate_columns(selection_frame), min_group_fraction=min_group_fraction, ) if len(columns) < 2: raise ValueError("mask clustering requires at least two non-constant missingness columns") selection_mask = _mask_feature_matrix(selection_frame, columns) calibration_mask = _mask_feature_matrix(calibration_frame, columns) test_mask = _mask_feature_matrix(test_frame, columns) n_components = min(10, selection_mask.shape[0], selection_mask.shape[1]) if n_components >= 1 and n_components < selection_mask.shape[1]: reducer = TruncatedSVD(n_components=n_components, random_state=random_state) selection_features = reducer.fit_transform(selection_mask) calibration_features = reducer.transform(calibration_mask) test_features = reducer.transform(test_mask) else: selection_features = selection_mask calibration_features = calibration_mask test_features = test_mask diagnostics: dict[int, dict[str, float | int]] = {} best_k = None best_score = None best_model = None for k in mask_cluster_k_grid: if k <= 1 or k > len(selection_features): diagnostics[k] = {"silhouette": float("nan"), "min_cluster_size": 0, "selection_gap": float("nan")} continue model = KMeans(n_clusters=k, random_state=random_state, n_init=10) selection_ids = model.fit_predict(selection_features) counts = np.bincount(selection_ids, minlength=k) min_cluster_size = int(counts.min()) if counts.size else 0 if min_cluster_size < min_selection_group_rows: diagnostics[k] = {"silhouette": float("nan"), "min_cluster_size": min_cluster_size, "selection_gap": float("nan")} continue silhouette = ( float(silhouette_score(selection_features, selection_ids)) if len(np.unique(selection_ids)) > 1 else float("nan") ) cluster_means = [] for cluster_id in range(k): cluster_mask = selection_ids == cluster_id cluster_means.append(float(selection_frame.loc[cluster_mask, "global_missing_rate"].mean())) selection_gap = float(max(cluster_means) - min(cluster_means)) if cluster_means else 0.0 diagnostics[k] = { "silhouette": silhouette, "min_cluster_size": min_cluster_size, "selection_gap": selection_gap, } score = (-k,) if best_score is None or score > best_score: best_score = score best_k = k best_model = model if best_model is None or best_k is None: raise ValueError("mask clustering could not find a feasible k") selection_ids = best_model.predict(selection_features) calibration_ids = best_model.predict(calibration_features) test_ids = best_model.predict(test_features) group_labels = {group_id: f"cluster_{group_id}" for group_id in range(best_k)} return MissingnessGroupingResult( selection_groups=_frame_groups_from_ids(selection_ids, group_labels), calibration_groups=_frame_groups_from_ids(calibration_ids, group_labels), test_groups=_frame_groups_from_ids(test_ids, group_labels), group_labels=group_labels, metadata={ "strategy": "mask_cluster", "selected_k": best_k, "group_source": "missingness_mask_cluster", "group_count": len(group_labels), "mask_cluster_diagnostics": diagnostics, }, ) __all__ = ["MissingnessGroupingResult", "build_missingness_groups"]