patient_leakage_detection/classes/alignment.py

"""
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": <reference-model patient id>,
                "co_cluster_agreement": <mean pairwise Jaccard>,
                "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