plotters-statistical 0.1.0

//! Residual plot series: residuals vs fitted values, with a zero-reference line
//! and an optional binned moving-average trend line.
//!
//! The trend is a moving average over binned x-ranges rather than a full LOESS
//! fit — enough to reveal whether residuals drift away from zero, which is the
//! diagnostic point, without the much larger scope of a local-regression
//! smoother.

use plotters::element::{Drawable, PointCollection};
use plotters::style::RGBColor;
use plotters_backend::{BackendCoord, DrawingBackend, DrawingErrorKind};

use crate::stats::StatsError;
use crate::style::{stroke_style, translucent_fill};

const DEFAULT_POINT_COLOR: RGBColor = RGBColor(0, 114, 178); // Okabe–Ito blue
const ZERO_LINE_GRAY: RGBColor = RGBColor(120, 120, 120);
const TREND_COLOR: RGBColor = RGBColor(213, 94, 0); // Okabe–Ito vermillion
const DEFAULT_BINS: usize = 12;

/// A residual plot as a drawable series (coordinate space `(f64, f64)` =
/// `(fitted, residual)`).
#[derive(Debug, Clone)]
pub struct ResidualPlot {
    // Layout: [scatter (n_points)] [zero-line anchors: (xmin,0),(xmax,0)]
    //         [trend points (n_trend)].
    points: Vec<(f64, f64)>,
    n_points: usize,
    n_trend: usize,
    show_trend: bool,
    marker_radius: u32,
    point_color: RGBColor,
    trend_width: u32,
}

impl ResidualPlot {
    /// Build from `fitted` (predicted) values and their `residuals`.
    pub fn from_residuals(fitted: &[f64], residuals: &[f64]) -> Result<Self, StatsError> {
        if fitted.len() != residuals.len() {
            return Err(StatsError::LengthMismatch {
                scores: fitted.len(),
                labels: residuals.len(),
            });
        }
        let pairs: Vec<(f64, f64)> = fitted
            .iter()
            .zip(residuals)
            .map(|(&x, &r)| (x, r))
            .filter(|(x, r)| x.is_finite() && r.is_finite())
            .collect();
        Self::from_pairs(pairs)
    }

    /// Build from `predicted` and `actual` values, computing residuals as
    /// `actual - predicted` internally.
    pub fn from_predictions(predicted: &[f64], actual: &[f64]) -> Result<Self, StatsError> {
        if predicted.len() != actual.len() {
            return Err(StatsError::LengthMismatch {
                scores: predicted.len(),
                labels: actual.len(),
            });
        }
        let residuals: Vec<f64> = predicted.iter().zip(actual).map(|(&p, &a)| a - p).collect();
        Self::from_residuals(predicted, &residuals)
    }

    fn from_pairs(pairs: Vec<(f64, f64)>) -> Result<Self, StatsError> {
        if pairs.is_empty() {
            return Err(StatsError::EmptyInput);
        }
        let n_points = pairs.len();
        let xmin = pairs.iter().map(|p| p.0).fold(f64::INFINITY, f64::min);
        let xmax = pairs.iter().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max);

        let trend = binned_trend(&pairs, DEFAULT_BINS);

        let mut points = pairs;
        points.push((xmin, 0.0));
        points.push((xmax, 0.0));
        let n_trend = trend.len();
        points.extend(trend);

        Ok(Self {
            points,
            n_points,
            n_trend,
            show_trend: false,
            marker_radius: 3,
            point_color: DEFAULT_POINT_COLOR,
            trend_width: 2,
        })
    }

    /// Enable/disable the binned moving-average trend line.
    pub fn trend(mut self, show: bool) -> Self {
        self.show_trend = show;
        self
    }

    /// Set the scatter marker radius in pixels.
    pub fn marker_radius(mut self, radius: u32) -> Self {
        self.marker_radius = radius;
        self
    }

    /// Set the scatter point color.
    pub fn color(mut self, color: RGBColor) -> Self {
        self.point_color = color;
        self
    }
}

impl<'a> PointCollection<'a, (f64, f64)> for &'a ResidualPlot {
    type Point = &'a (f64, f64);
    type IntoIter = &'a [(f64, f64)];
    fn point_iter(self) -> &'a [(f64, f64)] {
        &self.points
    }
}

impl<DB: DrawingBackend> Drawable<DB> for ResidualPlot {
    fn draw<I: Iterator<Item = BackendCoord>>(
        &self,
        points: I,
        backend: &mut DB,
        _parent_dim: (u32, u32),
    ) -> Result<(), DrawingErrorKind<DB::ErrorType>> {
        let pix: Vec<BackendCoord> = points.collect();
        if pix.len() < self.n_points + 2 {
            return Ok(());
        }
        // Zero reference line (drawn first, under the points).
        let anchor_l = pix[self.n_points];
        let anchor_r = pix[self.n_points + 1];
        backend.draw_line(anchor_l, anchor_r, &stroke_style(&ZERO_LINE_GRAY, 1))?;

        // Scatter.
        let marker = translucent_fill(&self.point_color, 0.7);
        for p in &pix[..self.n_points] {
            backend.draw_circle(*p, self.marker_radius, &marker, true)?;
        }

        // Trend line.
        if self.show_trend && self.n_trend >= 2 {
            let start = self.n_points + 2;
            let trend = &pix[start..start + self.n_trend];
            backend.draw_path(
                trend.iter().copied(),
                &stroke_style(&TREND_COLOR, self.trend_width),
            )?;
        }
        Ok(())
    }
}

/// Moving average of residuals over `bins` equal-width x-bins. Returns one
/// `(bin_center, mean_residual)` point per non-empty bin, in ascending x.
fn binned_trend(pairs: &[(f64, f64)], bins: usize) -> Vec<(f64, f64)> {
    let bins = bins.max(1);
    let xmin = pairs.iter().map(|p| p.0).fold(f64::INFINITY, f64::min);
    let xmax = pairs.iter().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max);
    if !(xmin.is_finite() && xmax.is_finite()) || xmax <= xmin {
        return Vec::new();
    }
    let width = (xmax - xmin) / bins as f64;
    let mut sums = vec![0.0_f64; bins];
    let mut counts = vec![0usize; bins];
    for &(x, r) in pairs {
        let mut idx = ((x - xmin) / width).floor() as usize;
        if idx >= bins {
            idx = bins - 1; // the max value lands in the last bin
        }
        sums[idx] += r;
        counts[idx] += 1;
    }
    (0..bins)
        .filter(|&b| counts[b] > 0)
        .map(|b| {
            let center = xmin + width * (b as f64 + 0.5);
            (center, sums[b] / counts[b] as f64)
        })
        .collect()
}