""" Align patient clusters across models using the Hungarian algorithm. Each model clusters independently, so its cluster-3 for benign may not correspond to another model's cluster-3. The Hungarian algorithm finds the optimal 1-to-1 mapping between cluster labels by maximizing image overlap. """ import numpy as np from scipy.optimize import linear_sum_assignment class PatientAligner: """Align patient clusters across multiple models. Parameters ---------- assignments : dict Nested dict: model_name -> {patient_id: [filenames, ...]} e.g. {"VGG16": {"benign_0": [...], ...}, "DenseNet121": {...}, ...} """ def __init__(self, assignments): self.assignments = assignments self.model_names = sorted(assignments.keys()) self.n_models = len(self.model_names) def align_to_reference(self, reference_model=None): """Align every model's clusters to a reference model. Parameters ---------- reference_model : str or None Model to use as reference. Defaults to the first model. Returns ------- mapping : dict model_name -> {model_patient_id: reference_patient_id} """ if reference_model is None: reference_model = self.model_names[0] print(f"\nAligning all models to reference: {reference_model}") mapping = {reference_model: {pid: pid for pid in self.assignments[reference_model]}} for model in self.model_names: if model == reference_model: continue print(f" {model} -> {reference_model} ...") mapping[model] = {} # Align class by class for prefix in self._class_prefixes(): mapping[model].update( self._align_class(model, reference_model, prefix) ) return mapping def build_consensus(self, mapping): """For each image, measure model agreement via co-clustering overlap. Instead of comparing arbitrary cluster IDs (even aligned ones), we ask: *which other images* does each model place in this image's cluster? High Jaccard overlap = models agree on patient membership regardless of what they call the cluster. Returns ------- consensus : dict image_filename -> { "consensus_pid": , "co_cluster_agreement": , "cluster_size": {model: n_images_in_cluster}, "raw_pids": {model: original_patient_id}, } """ # Build cluster lookup: model -> pid -> set of filenames clusters = {} for model, model_assignments in self.assignments.items(): clusters[model] = {pid: set(files) for pid, files in model_assignments.items()} # Per-image: find which cluster it's in per model image_to_cluster = {} for model, model_clusters in clusters.items(): for pid, files in model_clusters.items(): for f in files: image_to_cluster.setdefault(f, {})[model] = pid consensus = {} for img, model_pids in image_to_cluster.items(): # Co-cluster set for each model (images in same cluster, excluding self) co_sets = {} for model in self.model_names: pid = model_pids.get(model) if pid is not None: co_sets[model] = clusters[model][pid] - {img} else: co_sets[model] = set() # Pairwise Jaccard between all model pairs jaccards = [] for i, ma in enumerate(self.model_names): for j, mb in enumerate(self.model_names): if i < j: sa, sb = co_sets[ma], co_sets[mb] union = len(sa | sb) jac = len(sa & sb) / union if union > 0 else 1.0 jaccards.append(jac) mean_jaccard = np.mean(jaccards) if jaccards else 1.0 # Consensus PID via aligned majority vote aligned_votes = {} for model, pid in model_pids.items(): aligned_votes[model] = mapping.get(model, {}).get(pid, pid) pid_counts = {} for pid in aligned_votes.values(): pid_counts[pid] = pid_counts.get(pid, 0) + 1 consensus_pid = max(pid_counts, key=pid_counts.get) consensus[img] = { "consensus_pid": consensus_pid, "co_cluster_agreement": round(mean_jaccard, 3), "cluster_size": {m: len(co_sets[m]) + 1 for m in self.model_names}, "raw_pids": model_pids, } return consensus # ------------------------------------------------------------------ # Internals # ------------------------------------------------------------------ def _class_prefixes(self): """Return the set of class prefixes (benign, malig, normal).""" prefixes = set() for assignments in self.assignments.values(): for pid in assignments: prefixes.add(pid.rsplit("_", 1)[0]) return sorted(prefixes) def _align_class(self, model_a, model_b, prefix): """Hungarian-align patient IDs for one class between two models. Returns dict: model_a_patient_id -> model_b_patient_id """ # Gather patient IDs for this class prefix pids_a = sorted([p for p in self.assignments[model_a] if p.startswith(prefix)]) pids_b = sorted([p for p in self.assignments[model_b] if p.startswith(prefix)]) n_a, n_b = len(pids_a), len(pids_b) if n_a == 0 or n_b == 0: return {} # Build overlap cost matrix (negate overlap → minimize = maximize overlap) cost = np.zeros((n_a, n_b)) for i, pa in enumerate(pids_a): files_a = set(self.assignments[model_a][pa]) for j, pb in enumerate(pids_b): files_b = set(self.assignments[model_b][pb]) overlap = len(files_a & files_b) cost[i, j] = -overlap # Hungarian *minimizes* cost row_ind, col_ind = linear_sum_assignment(cost) # Build mapping mapping = {} zero_matches = 0 missed_opportunities = 0 for i, j in zip(row_ind, col_ind): overlap = -cost[i, j] total_a = len(self.assignments[model_a][pids_a[i]]) # Best possible match for this cluster (even if Hungarian didn't pick it) best_overlap = int(-cost[i].min()) # most negative → best match best_j = int(np.argmin(cost[i])) mapping[pids_a[i]] = pids_b[j] if overlap == 0: zero_matches += 1 if best_overlap > 0: missed_opportunities += 1 best_pid = pids_b[best_j] print(f" LOW-OVERLAP: {pids_a[i]} ({total_a} imgs) matched to " f"{pids_b[j]} (0 overlap) — best would be {best_pid} " f"({best_overlap}/{total_a} overlap)") else: print(f" UNMATCHED: {pids_a[i]} ({total_a} imgs) — " f"no cluster in {model_b} shares ANY images with this group") if zero_matches: print(f" {zero_matches} zero-overlap matches ({missed_opportunities} " f"forced by global optimum, {zero_matches - missed_opportunities} " f"truly unmatched)") return mapping