plotters-statistical 0.1.0

//! ROC curve: threshold sweep producing (FPR, TPR) points, and AUC via
//! trapezoidal integration.

use super::StatsError;

/// A single operating point on an ROC curve.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RocPoint {
    /// False positive rate, `FP / N`.
    pub fpr: f64,
    /// True positive rate (recall / sensitivity), `TP / P`.
    pub tpr: f64,
    /// The score threshold at which this point is reached (predictions with
    /// score `>= threshold` are called positive). The leading `(0, 0)` point
    /// carries `+inf`.
    pub threshold: f64,
}

/// An ROC curve: its operating points (from `(0,0)` to `(1,1)`) and the area
/// under it.
#[derive(Debug, Clone, PartialEq)]
pub struct RocCurveData {
    /// Operating points, ascending in FPR.
    pub points: Vec<RocPoint>,
    /// Area under the ROC curve, in `[0, 1]`.
    pub auc: f64,
}

/// Trapezoidal integration of `y` over `x` for a sequence of `(x, y)` points
/// assumed sorted by ascending `x`.
///
/// Exposed because AUC is just this applied to `(fpr, tpr)`; callers with their
/// own point sets can reuse it directly.
pub fn auc_trapezoid(points: &[(f64, f64)]) -> f64 {
    points
        .windows(2)
        .map(|w| {
            let (x1, y1) = w[0];
            let (x2, y2) = w[1];
            (x2 - x1) * (y1 + y2) / 2.0
        })
        .sum()
}

/// Compute an ROC curve from predicted `scores` and true binary `labels`
/// (`true` = positive class).
///
/// Higher scores are treated as "more positive". The sweep visits each distinct
/// score as a threshold, emitting one point per distinct score (ties handled
/// together), and always begins at `(0, 0)`.
///
/// # Errors
///
/// * [`StatsError::LengthMismatch`] if `scores` and `labels` differ in length.
/// * [`StatsError::EmptyInput`] if empty.
/// * [`StatsError::NoPositiveLabels`] / [`StatsError::NoNegativeLabels`] if a
///   class is absent, since TPR or FPR would be `0/0`.
pub fn roc_curve(scores: &[f64], labels: &[bool]) -> Result<RocCurveData, StatsError> {
    if scores.len() != labels.len() {
        return Err(StatsError::LengthMismatch {
            scores: scores.len(),
            labels: labels.len(),
        });
    }
    if scores.is_empty() {
        return Err(StatsError::EmptyInput);
    }
    let total_p = labels.iter().filter(|&&l| l).count();
    let total_n = labels.len() - total_p;
    if total_p == 0 {
        return Err(StatsError::NoPositiveLabels);
    }
    if total_n == 0 {
        return Err(StatsError::NoNegativeLabels);
    }

    // Sort indices by descending score.
    let mut order: Vec<usize> = (0..scores.len()).collect();
    order.sort_by(|&a, &b| {
        scores[b]
            .partial_cmp(&scores[a])
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    let (p, n) = (total_p as f64, total_n as f64);
    let mut points = Vec::with_capacity(order.len() + 1);
    points.push(RocPoint {
        fpr: 0.0,
        tpr: 0.0,
        threshold: f64::INFINITY,
    });

    let mut tp = 0usize;
    let mut fp = 0usize;
    let mut i = 0usize;
    while i < order.len() {
        let score = scores[order[i]];
        // Consume every instance sharing this score before emitting a point,
        // so tied scores can't be split by an arbitrary threshold.
        while i < order.len() && scores[order[i]] == score {
            if labels[order[i]] {
                tp += 1;
            } else {
                fp += 1;
            }
            i += 1;
        }
        points.push(RocPoint {
            fpr: fp as f64 / n,
            tpr: tp as f64 / p,
            threshold: score,
        });
    }

    let auc = auc_trapezoid(&points.iter().map(|pt| (pt.fpr, pt.tpr)).collect::<Vec<_>>());

    Ok(RocCurveData { points, auc })
}