Blame for scripts/validation/analyse_sensitivity_noB6.py

"""
Sensitivity analysis with known malignancy removed: no BI-RADS 6, no staging examinations.

    python scripts/validation/analyse_sensitivity_noB6.py <oof_dir> <metadata_file>

BI-RADS 6 means a biopsy proven carcinoma is already known, and a staging examination is performed
because a carcinoma has been diagnosed. Neither can contribute to a *detection* claim: the model is
not finding anything that was not already known. This script removes those breasts and repeats the
primary analysis on the models trained on the full data, answering whether the non-inferiority
claim survives.

Output is one row per model setting with

    N breasts, AUC (95% CI), sensitivity at 90% specificity (95% CI),
    specificity at 90% sensitivity (95% CI),
    delta AUC (AP - FDP) with the 97.5% Bonferroni adjusted interval, and the verdict at the margin

Everything is computed the way the primary analysis computes it: metrics within each outer fold,
combined with fixed weights, intervals from a bootstrap that resamples whole women within fold. The
delta is paired, both protocols scored on the same resampled women, and its interval is the 97.5%
one, which is the Bonferroni adjustment for the two architectures. The delta and the verdict belong
to the architecture, so the two rows of one architecture repeat them.

Breasts whose BI-RADS or indication is unknown cannot be checked for known malignancy. --unknown
decides what happens to them; the report always states how many there are and what the numbers look
like either way, because on this cohort that choice moves the case mix substantially.
"""

import argparse
from pathlib import Path

import numpy as np
import pandas as pd

from auto_detect_breast_mri.config import (normalise_indication_code, require_indication_labels,
                                           resolve_path)
from auto_detect_breast_mri.evaluation.cluster_bootstrap import (FoldClusterBootstrap,
                                                                 percentile_interval,
                                                                 weighted_fold_difference)
from auto_detect_breast_mri.evaluation.oof_table import (ARCHITECTURE_LABELS, PROTOCOL_LABELS,
                                                         available_fractions, load_patient_map,
                                                         load_side_map)
from auto_detect_breast_mri.evaluation.subgroups import bootstrap_metrics

DEFAULT_OPERATING_TARGET = 0.9
DEFAULT_MARGIN = 0.03
EXCLUDED_BIRADS = '6'
EXCLUDED_INDICATION = 'Staging'
UNKNOWN = 'unknown'
MODEL_ORDER = [('resnet18', 'abrv'), ('resnet18', 'full'),
               ('resnet50', 'abrv'), ('resnet50', 'full')]


def model_name(architecture, protocol):
    return (f"{ARCHITECTURE_LABELS.get(architecture, architecture)} "
            f"{PROTOCOL_LABELS.get(protocol, protocol)}")


def read_predictions(oof_dirs, metadata_file, fraction=None):
    """
    Wide table, one row per breast, with a score column per model. Full data models by default.

    :param oof_dirs: one folder, or several when the architectures were predicted into separate
        output roots
    :return: (DataFrame, dict describing what was read)
    """
    from auto_detect_breast_mri.evaluation.oof_table import FILE_PATTERN

    if isinstance(oof_dirs, (str, Path)):
        oof_dirs = [oof_dirs]
    per_fraction = {}
    for oof_dir in oof_dirs:
        if not Path(oof_dir).is_dir():
            raise FileNotFoundError(f"{oof_dir} is not a directory.")
        for value, paths in available_fractions(oof_dir).items():
            per_fraction.setdefault(value, []).extend(paths)
    if not per_fraction:
        raise FileNotFoundError(f"No oof_*_fold*.csv in {', '.join(map(str, oof_dirs))}. "
                                f"Run predict_oof.py first.")
    chosen = max(per_fraction) if fraction is None else float(fraction)
    if chosen not in per_fraction:
        raise ValueError(f"No predictions for fraction {chosen:g}; the folders hold "
                         f"{', '.join(f'{value:g}' for value in sorted(per_fraction))}.")
    if chosen != 1.0:
        print(f"  NOTE: analysing fraction {chosen:g}, not the full data models.")

    frames = []
    for path in sorted(per_fraction[chosen]):
        parsed = FILE_PATTERN.match(Path(path).name)
        frame = pd.read_csv(path, dtype={'examination_id': str, 'side': str})
        frame['architecture'] = parsed.group('architecture')
        frames.append(frame)
    long_table = pd.concat(frames, ignore_index=True)
    long_table['side'] = long_table['side'].str.lower()
    long_table['column'] = long_table['architecture'] + '_' + long_table['protocol']

    keys = ['outer_fold', 'examination_id', 'side', 'label']
    wide = long_table.pivot_table(index=keys, columns='column', values='score').reset_index()
    wide.columns.name = None

    patient_map = load_patient_map(metadata_file, wide['examination_id'].unique())
    wide['patient_id'] = wide['examination_id'].map(patient_map)
    report = {'fraction': chosen, 'breasts_total': len(wide),
              'malignant_total': int(wide['label'].sum()),
              'models': [column for column in wide.columns
                         if any(column == f'{a}_{p}' for a, p in MODEL_ORDER)]}
    return wide, report


def attach_status(table, metadata_file):
    """
    Add the BI-RADS category and the indication group of every breast, both per side.

    :return: (table with 'birads' and 'indication_group' columns, notes dict)
    """
    table = table.copy()
    identifiers = table['examination_id'].unique()
    pairs = list(zip(table['examination_id'], table['side']))
    notes = {}

    birads, birads_report = load_side_map(metadata_file, 'birads', identifiers,
                                          normalise=normalise_indication_code)
    if birads is None:
        raise ValueError(f"No BI-RADS columns in the metadata "
                         f"({', '.join(birads_report['missing_columns'])}); this analysis needs "
                         f"them. Configure columns.birads_left / columns.birads_right.")
    table['birads'] = [birads.get(pair, UNKNOWN) for pair in pairs]

    indication, _ = load_side_map(metadata_file, 'indication', identifiers,
                                  normalise=normalise_indication_code)
    labels = require_indication_labels()
    table['indication_group'] = [
        labels.get(indication.get(pair), labels.get(UNKNOWN, UNKNOWN)) if indication else UNKNOWN
        for pair in pairs]

    notes['birads_unknown'] = int((table['birads'] == UNKNOWN).sum())
    notes['indication_unknown'] = int((table['indication_group'] == UNKNOWN).sum())
    return table, notes


def select(table, unknown='keep'):
    """
    Drop the breasts that represent already known malignancy.

    :param unknown: 'keep' leaves breasts whose status cannot be checked in the analysis,
        'exclude' removes them as well
    :return: (kept table, dict of counts)
    """
    known_malignancy = ((table['birads'] == EXCLUDED_BIRADS)
                        | (table['indication_group'] == EXCLUDED_INDICATION))
    unverifiable = (table['birads'] == UNKNOWN) & (table['indication_group'] == UNKNOWN)

    counts = {
        'birads6': int((table['birads'] == EXCLUDED_BIRADS).sum()),
        'staging': int((table['indication_group'] == EXCLUDED_INDICATION).sum()),
        'either': int(known_malignancy.sum()),
        'unverifiable': int(unverifiable.sum()),
        'unverifiable_malignant': int(table.loc[unverifiable, 'label'].sum()),
    }
    kept = table[~known_malignancy]
    if unknown == 'exclude':
        kept = kept[~unverifiable]
    counts['kept'] = len(kept)
    counts['kept_malignant'] = int(kept['label'].sum())
    return kept.reset_index(drop=True), counts


def paired_delta(table, ap_column, fdp_column, replications, weighting, seed, stratify,
                 confidence=0.975):
    """
    Paired AUC(AP) - AUC(FDP), both protocols scored on the same resampled women.

    :param confidence: 0.975 is the Bonferroni adjusted level for two architectures
    :return: dict with the observed delta and its interval
    """
    bootstrap = FoldClusterBootstrap(table['outer_fold'].to_numpy(),
                                     table['patient_id'].to_numpy(),
                                     table['label'].to_numpy(), stratify_by_outcome=stratify)
    weights = bootstrap.fold_weights(weighting)
    labels = table['label'].to_numpy()
    ap_scores, fdp_scores = table[ap_column].to_numpy(), table[fdp_column].to_numpy()

    observed, _ = weighted_fold_difference(bootstrap.observed_rows(), labels, ap_scores,
                                           fdp_scores, weights)
    rng = np.random.default_rng(seed)
    replicated = np.full(replications, np.nan)
    for replication in range(replications):
        drawn = bootstrap.resample_rows(rng)
        replicated[replication], _ = weighted_fold_difference(drawn, labels, ap_scores,
                                                              fdp_scores, weights)
    low, high = percentile_interval(replicated, confidence)
    return {'delta': observed, 'ci_low': low, 'ci_high': high,
            'degenerate': int(np.isnan(replicated).sum())}


def analyse(table, target, margin, replications, weighting, seed, stratify, decimals=3):
    """
    One row per model setting, with the architecture level delta repeated on both of its rows.

    :return: (DataFrame of strings, list of notes, list of raw result dicts)
    """
    def interval(value, low, high):
        if np.isnan(value):
            return '-'
        return f"{value:.{decimals}f} ({low:.{decimals}f}-{high:.{decimals}f})"

    deltas, notes, raw = {}, [], []
    architectures = sorted({architecture for architecture, _ in MODEL_ORDER
                            if f'{architecture}_abrv' in table.columns
                            and f'{architecture}_full' in table.columns})
    for architecture in architectures:
        print(f"  paired delta for {architecture} ...", flush=True)
        deltas[architecture] = paired_delta(table, f'{architecture}_abrv', f'{architecture}_full',
                                            replications, weighting, seed, stratify)
        if deltas[architecture]['degenerate']:
            notes.append(f"{architecture}: {deltas[architecture]['degenerate']} of {replications} "
                         f"replications had an undefined AUC in every fold and were dropped")

    rows = []
    for architecture, protocol in MODEL_ORDER:
        column = f'{architecture}_{protocol}'
        if column not in table.columns:
            continue
        print(f"  {column} ...", flush=True)
        statistics = bootstrap_metrics(table['outer_fold'].to_numpy(),
                                       table['patient_id'].to_numpy(), table['label'].to_numpy(),
                                       table[column].to_numpy(), target=target,
                                       replications=replications, weighting=weighting, seed=seed,
                                       stratify=stratify)
        delta = deltas.get(architecture)
        met = delta is not None and not np.isnan(delta['ci_low']) and delta['ci_low'] > -margin
        rows.append({
            'Model setting': model_name(architecture, protocol),
            'N breasts (without BI-RADS 6)': len(table),
            'AUC (95% CI)': interval(statistics['auc']['value'], statistics['auc']['ci_low'],
                                     statistics['auc']['ci_high']),
            f'Sensitivity @ {target:.0%} specificity (95% CI)':
                interval(statistics['sens_at_spec']['value'],
                         statistics['sens_at_spec']['ci_low'],
                         statistics['sens_at_spec']['ci_high']),
            f'Specificity @ {target:.0%} sensitivity (95% CI)':
                interval(statistics['spec_at_sens']['value'],
                         statistics['spec_at_sens']['ci_low'],
                         statistics['spec_at_sens']['ci_high']),
            'ΔAUC (AP−FDP), 97.5% Bonferroni-adjusted CI':
                interval(delta['delta'], delta['ci_low'], delta['ci_high']) if delta else '-',
            f'Non-inferiority verdict (margin {margin})': ('met' if met else 'not met')
                if delta else '-',
        })
        raw.append({'architecture': architecture, 'protocol': protocol, 'breasts': len(table),
                    **{f'{metric}_{field}': value for metric, values in statistics.items()
                       for field, value in values.items()},
                    **({f'delta_{field}': value for field, value in delta.items()} if delta else {})})
    return pd.DataFrame(rows), notes, raw


def format_report(frame, counts, status_notes, report, notes, target, margin, replications,
                  unknown):
    lines = []
    add = lines.append
    add("=" * 110)
    add("Sensitivity analysis: BI-RADS 6 and staging examinations removed")
    add("=" * 110)
    add("")
    add(f"{'training fraction':<34}{report['fraction']:g}")
    add(f"{'breasts before exclusion':<34}{report['breasts_total']} "
        f"({report['malignant_total']} malignant)")
    add(f"{'BI-RADS 6':<34}{counts['birads6']}")
    add(f"{'staging indication':<34}{counts['staging']}")
    add(f"{'either (removed)':<34}{counts['either']}")
    add(f"{'status unknown':<34}{counts['unverifiable']} "
        f"({counts['unverifiable_malignant']} malignant) -> {unknown}")
    add(f"{'breasts analysed':<34}{counts['kept']} ({counts['kept_malignant']} malignant)")
    add(f"{'bootstrap replications':<34}{replications}")
    add("")
    add("-" * 110)
    add("Result")
    add("-" * 110)
    widths = {column: max(len(str(column)), *(len(str(value)) for value in frame[column]))
              for column in frame.columns}
    add('  '.join(str(column).ljust(widths[column]) for column in frame.columns))
    add('  '.join('-' * widths[column] for column in frame.columns))
    for _, row in frame.iterrows():
        add('  '.join(str(row[column]).ljust(widths[column]) for column in frame.columns))
    add("")
    for note in notes:
        add(f"NOTE: {note}")
    add("")
    add("-" * 110)
    add("How to read this")
    add("-" * 110)
    add("Delta and verdict belong to the architecture, not to a single protocol, so the two rows")
    add("of one architecture carry the same pair. Non-inferiority is met when the lower bound of")
    add(f"the 97.5% interval, the Bonferroni adjustment for two architectures, exceeds -{margin}.")
    add("")
    add("The two operating points take their threshold from the fold they are scored on and are")
    add("therefore optimistic; the AUC is threshold free. Intervals resample whole women within")
    add("each outer fold, so both breasts and repeated examinations of one woman move together,")
    add("and are conditional on the trained models: they do not carry training variability.")
    add("")
    add(f"BI-RADS unknown for {status_notes['birads_unknown']} breasts, indication unknown for "
        f"{status_notes['indication_unknown']}. Those breasts cannot be checked for already known")
    add("malignancy; --unknown decides whether they stay in. Rerun with the other setting to see")
    add("how much that choice matters on this cohort.")
    return "\n".join(lines)


def build_parser():
    parser = argparse.ArgumentParser(prog="analyse_sensitivity_noB6", description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("oof_dirs", nargs='+', metavar="OOF_DIR",
                        help="Folder(s) holding oof_<arch>_fold<k>[_frac=<f>].csv. Give several "
                             "when the architectures were predicted into separate output roots.")
    parser.add_argument("metadata_file", help="Metadata export with BI-RADS and indication.")
    parser.add_argument("-o", "--output_path", default=None)
    parser.add_argument("--fraction", type=float, default=None,
                        help="Training fraction. Default: the largest present, i.e. the full "
                             "data models this analysis is meant for.")
    parser.add_argument("--unknown", choices=['keep', 'exclude'], default='keep',
                        help="What to do with breasts whose BI-RADS and indication are both "
                             "unknown, so known malignancy cannot be ruled out. Default: keep")
    parser.add_argument("--operating_point", type=float, default=DEFAULT_OPERATING_TARGET)
    parser.add_argument("--margin", type=float, default=DEFAULT_MARGIN,
                        help="Non-inferiority margin. Default: %(default)s")
    parser.add_argument("-r", "--replications", type=int, default=2000)
    parser.add_argument("-w", "--weighting", choices=["pairs", "cases", "equal"], default="pairs")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--no_stratify", action="store_true")
    parser.add_argument("-d", "--decimals", type=int, default=3)
    return parser


def main():
    args = build_parser().parse_args()

    output_dir = Path(resolve_path(args.output_path, "output_root", "output folder")) / "sensitivity_noB6"
    output_dir.mkdir(parents=True, exist_ok=True)

    table, report = read_predictions(args.oof_dirs, args.metadata_file, args.fraction)
    table, status_notes = attach_status(table, args.metadata_file)
    kept, counts = select(table, args.unknown)
    if kept.empty:
        raise ValueError("Nothing left after the exclusions.")
    missing_patient = int(kept['patient_id'].isna().sum())
    if missing_patient:
        print(f"  {missing_patient} breasts without a patient ID are dropped from the clustering")
        kept = kept[kept['patient_id'].notna()].reset_index(drop=True)
        counts['kept'], counts['kept_malignant'] = len(kept), int(kept['label'].sum())

    print(f"Bootstrapping with {args.replications} replications ...")
    frame, notes, raw = analyse(kept, args.operating_point, args.margin, args.replications,
                                args.weighting, args.seed, not args.no_stratify, args.decimals)

    csv_path = output_dir / "sensitivity_noB6_table.csv"
    frame.to_csv(csv_path, index=False)
    pd.DataFrame(raw).to_csv(output_dir / "sensitivity_noB6_raw.csv", index=False)
    text = format_report(frame, counts, status_notes, report, notes, args.operating_point,
                         args.margin, args.replications, args.unknown)
    (output_dir / "sensitivity_noB6_report.txt").write_text(text)
    print()
    print(text)
    print(f"\nWrote {csv_path}")


if __name__ == '__main__':
    main()