burn-train 0.22.1

use std::sync::Arc;

use super::state::{FormatOptions, NumericMetricState};
use super::{MetricMetadata, SerializedEntry};
use crate::metric::{
    Metric, MetricAttributes, MetricName, Numeric, NumericAttributes, NumericEntry,
};
use burn_core::tensor::{Int, Tensor, TensorReadError};

/// The Top-K accuracy metric.
///
/// For K=1 (the default), this is equivalent to the [accuracy metric](`super::acc::AccuracyMetric`).
///
/// Updating the metric panics if K exceeds the number of output classes.
#[derive(Clone)]
pub struct TopKAccuracyMetric {
    name: Arc<String>,
    k: usize,
    state: NumericMetricState,
    /// If specified, targets equal to this value will be considered padding and will not count
    /// towards the metric
    pad_token: Option<usize>,
}

/// The [top-k accuracy metric](TopKAccuracyMetric) input type.
#[derive(new)]
pub struct TopKAccuracyInput {
    /// The outputs (batch_size, num_classes)
    outputs: Tensor<2>,
    /// The labels (batch_size)
    targets: Tensor<1, Int>,
}

impl Default for TopKAccuracyMetric {
    fn default() -> Self {
        Self::new(1)
    }
}

impl TopKAccuracyMetric {
    /// Creates the metric.
    ///
    /// # Panics
    ///
    /// Panics if `k` is zero.
    pub fn new(k: usize) -> Self {
        assert!(
            k > 0,
            "TopKAccuracyMetric requires k to be greater than zero"
        );

        Self {
            name: Arc::new(format!("Top-K Accuracy @ TopK({})", k)),
            k,
            state: Default::default(),
            pad_token: None,
        }
    }

    /// Sets the pad token.
    pub fn with_pad_token(mut self, index: usize) -> Self {
        self.pad_token = Some(index);
        self
    }
}

impl Metric for TopKAccuracyMetric {
    type Input = TopKAccuracyInput;

    fn update(
        &mut self,
        input: &TopKAccuracyInput,
        _metadata: &MetricMetadata,
    ) -> Result<SerializedEntry, TensorReadError> {
        let [batch_size, n_classes] = input.outputs.dims();
        assert!(
            self.k <= n_classes,
            "TopKAccuracyMetric requires k ({}) to be no greater than the number of classes ({})",
            self.k,
            n_classes
        );

        let targets = input.targets.clone();

        let outputs = input
            .outputs
            .clone()
            .argsort_descending(1)
            .narrow(1, 0, self.k)
            .reshape([batch_size, self.k]);

        let (targets, num_pad) = match self.pad_token {
            Some(pad_token) => {
                // we ignore the samples where the target is equal to the pad token
                let mask = targets.clone().equal_scalar(pad_token as i64);
                let num_pad = mask.clone().int().sum().try_into_scalar::<i64>()? as usize;
                (targets.clone().mask_fill(mask, -1_i64), num_pad)
            }
            None => (targets.clone(), 0),
        };
        let valid_count = batch_size - num_pad;

        let accuracy = targets
            .reshape([batch_size, 1])
            .repeat_dim(1, self.k)
            .equal(outputs)
            .int()
            .sum()
            .try_into_scalar::<f64>()?
            / valid_count as f64;

        self.state.update(100.0 * accuracy, valid_count);
        Ok(self
            .state
            .compute_update(FormatOptions::new(self.name()).unit("%").precision(2)))
    }

    fn compute(&mut self) -> Result<SerializedEntry, TensorReadError> {
        Ok(self
            .state
            .compute_final(FormatOptions::new(self.name()).unit("%").precision(2)))
    }
    fn clear(&mut self) {
        self.state.reset()
    }

    fn name(&self) -> MetricName {
        self.name.clone()
    }

    fn attributes(&self) -> MetricAttributes {
        NumericAttributes {
            unit: Some("%".to_string()),
            higher_is_better: true,
        }
        .into()
    }
}

impl Numeric for TopKAccuracyMetric {
    fn value(&self) -> Option<NumericEntry> {
        Some(self.state.current_value())
    }

    fn running_value(&self) -> Option<NumericEntry> {
        Some(self.state.running_value())
    }

    fn final_value(&self) -> NumericEntry {
        self.state.final_value()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    #[should_panic(expected = "TopKAccuracyMetric requires k to be greater than zero")]
    fn test_zero_k_is_rejected() {
        TopKAccuracyMetric::new(0);
    }

    #[test]
    #[should_panic(
        expected = "TopKAccuracyMetric requires k (4) to be no greater than the number of classes (3)"
    )]
    fn test_k_greater_than_class_count_is_rejected() {
        let device = Default::default();
        let mut metric = TopKAccuracyMetric::new(4);
        let input = TopKAccuracyInput::new(
            Tensor::from_data([[0.1, 0.2, 0.7]], &device),
            Tensor::from_data([2], &device),
        );

        metric.update(&input, &MetricMetadata::fake());
    }

    #[test]
    fn test_k_equal_to_class_count() {
        let device = Default::default();
        let mut metric = TopKAccuracyMetric::new(3);
        let input = TopKAccuracyInput::new(
            Tensor::from_data([[0.1, 0.2, 0.7]; 3], &device),
            Tensor::from_data([0, 1, 2], &device),
        );

        metric.update(&input, &MetricMetadata::fake());
        assert_eq!(100.0, metric.value().unwrap().current());
    }

    #[test]
    fn test_default_uses_top_one() {
        let device = Default::default();
        let mut metric = TopKAccuracyMetric::default();
        let input = TopKAccuracyInput::new(
            Tensor::from_data([[0.1, 0.2, 0.7]; 2], &device),
            Tensor::from_data([2, 1], &device),
        );

        assert_eq!(metric.name(), TopKAccuracyMetric::new(1).name());
        metric.update(&input, &MetricMetadata::fake());
        assert_eq!(50.0, metric.value().unwrap().current());
    }

    #[test]
    fn test_accuracy_without_padding() {
        let device = Default::default();
        let mut metric = TopKAccuracyMetric::new(2);
        let input = TopKAccuracyInput::new(
            Tensor::from_data(
                [
                    [0.0, 0.2, 0.8], // 2, 1
                    [1.0, 2.0, 0.5], // 1, 0
                    [0.4, 0.1, 0.2], // 0, 2
                    [0.6, 0.7, 0.2], // 1, 0
                ],
                &device,
            ),
            Tensor::from_data([2, 2, 1, 1], &device),
        );

        let _entry = metric.update(&input, &MetricMetadata::fake()).unwrap();
        assert_eq!(50.0, metric.value().unwrap().current());
    }

    #[test]
    fn test_accuracy_with_padding() {
        let device = Default::default();
        let mut metric = TopKAccuracyMetric::new(2).with_pad_token(3);
        let input = TopKAccuracyInput::new(
            Tensor::from_data(
                [
                    [0.0, 0.2, 0.8, 0.0], // 2, 1
                    [1.0, 2.0, 0.5, 0.0], // 1, 0
                    [0.4, 0.1, 0.2, 0.0], // 0, 2
                    [0.6, 0.7, 0.2, 0.0], // 1, 0
                    [0.0, 0.1, 0.2, 5.0], // Predicted padding should not count
                    [0.0, 0.1, 0.2, 0.0], // Error on padding should not count
                    [0.6, 0.0, 0.2, 0.0], // Error on padding should not count
                ],
                &device,
            ),
            Tensor::from_data([2, 2, 1, 1, 3, 3, 3], &device),
        );

        let _entry = metric.update(&input, &MetricMetadata::fake()).unwrap();
        assert_eq!(50.0, metric.value().unwrap().current());
    }

    #[test]
    fn test_accuracy_epoch_aggregation_excludes_padding() {
        let device = Default::default();
        let mut metric = TopKAccuracyMetric::new(1).with_pad_token(2);

        // One valid, correct sample and three padding samples.
        metric
            .update(
                &TopKAccuracyInput::new(
                    Tensor::from_data([[0.9, 0.1], [0.9, 0.1], [0.9, 0.1], [0.9, 0.1]], &device),
                    Tensor::from_data([0, 2, 2, 2], &device),
                ),
                &MetricMetadata::fake(),
            )
            .unwrap();
        // Four valid, incorrect samples.
        metric
            .update(
                &TopKAccuracyInput::new(
                    Tensor::from_data([[0.9, 0.1]; 4], &device),
                    Tensor::from_data([1, 1, 1, 1], &device),
                ),
                &MetricMetadata::fake(),
            )
            .unwrap();

        // One correct prediction out of five valid samples.
        assert_eq!(20.0, metric.final_value().current());
    }

    #[test]
    fn test_parameterized_unique_name() {
        let metric_a = TopKAccuracyMetric::new(2);
        let metric_b = TopKAccuracyMetric::new(1);
        let metric_c = TopKAccuracyMetric::new(2);

        assert_ne!(metric_a.name(), metric_b.name());
        assert_eq!(metric_a.name(), metric_c.name());
    }
}