esoc-chart 0.1.0

High-level charting API built on esoc-gfx — matplotlib-equivalent for Rust
Documentation
// SPDX-License-Identifier: MIT OR Apache-2.0
//! Confusion matrix visualization.

use crate::express::heatmap;
use crate::grammar::chart::Chart;
use crate::new_theme::NewTheme;

/// Extension trait for creating confusion matrix figures.
pub trait ConfusionMatrixExt {
    /// Create a heatmap chart from this confusion matrix.
    fn figure(&self) -> Chart;

    /// Create a heatmap chart from this confusion matrix with a custom theme.
    fn figure_with_theme(&self, theme: NewTheme) -> Chart;
}

impl ConfusionMatrixExt for scry_learn::metrics::ConfusionMatrix {
    fn figure(&self) -> Chart {
        self.figure_with_theme(NewTheme::default())
    }

    fn figure_with_theme(&self, theme: NewTheme) -> Chart {
        let data: Vec<Vec<f64>> = self
            .matrix
            .iter()
            .map(|row| row.iter().map(|&v| v as f64).collect())
            .collect();

        heatmap(data)
            .annotate()
            .row_labels(self.labels.clone())
            .col_labels(self.labels.clone())
            .title("Confusion Matrix")
            .x_label("Predicted")
            .y_label("True")
            .theme(theme)
            .size(600.0, 600.0)
            .build()
    }
}

/// Create a confusion matrix figure.
pub fn confusion_matrix_figure(cm: &scry_learn::metrics::ConfusionMatrix) -> Chart {
    cm.figure()
}