Source code for compml.frequentist.tests.non_parametrics.friedman_aligned_ranks

"""Friedman aligned ranks test."""

from typing import Any

import numpy as np
from scipy.stats import chi2, rankdata

from compml.evaluation_scores import EvaluationScores
from compml.frequentist.tests.base import BaseFrequentistTest




class FriedmanAlignedRanksTest(BaseFrequentistTest[EvaluationScores]):
    """Friedman aligned ranks test for repeated measurements.

    The null hypothesis is that the mean scores are equal across all algorithms.
    """

    def __init__(self, alpha: float = 0.05, minimize: bool = False):
        """Initialize the Friedman aligned ranks test.

        Parameters
        ----------
        alpha : float
            Significance level.
        minimize : bool
            Whether to minimize the metric (lower scores are better).
        """
        super().__init__(alpha=alpha)
        self.__validate_arguments(minimize)
        self.minimize = minimize



    def compute_statistic(self, scores: EvaluationScores) -> float:
        """Compute the Friedman aligned ranks test statistic.

        Parameters
        ----------
        scores : EvaluationScores
            Evaluation scores container.

        Returns
        -------
        float
            Friedman aligned ranks test statistic.
        """
        x = scores.compute_means()
        n, k = x.shape
        aligned = x - x.mean(axis=1, keepdims=True)
        flat = aligned.ravel()
        if self.minimize:
            r = rankdata(flat, method="average")
        else:
            r = rankdata(-flat, method="average")
        ranks = r.reshape(n, k)

        sum_sq_col_sums = float(np.sum(ranks.sum(axis=0) ** 2))
        sum_sq_row_sums = float(np.sum(ranks.sum(axis=1) ** 2))
        nk = n * k
        numer = (k - 1) * (sum_sq_col_sums - (k * (n**2) / 4.0) * (nk + 1) ** 2)
        denom = (nk * (nk + 1) * (2 * nk + 1)) / 6.0 - (sum_sq_row_sums / k)

        if abs(denom) < 1e-15:
            return 0.0
        return numer / denom




    def compute_p_value(self, scores: EvaluationScores, statistic: float) -> float:
        """Compute p-value from the Friedman aligned ranks test statistic.

        Parameters
        ----------
        scores : EvaluationScores
            Evaluation scores container.
        statistic : float
            Friedman aligned ranks test statistic.

        Returns
        -------
        float
            p-value computed from the Friedman aligned ranks test statistic.
        """
        df = scores.n_algorithms - 1
        return float(chi2.sf(statistic, df))




    def compute_c_value(self, scores: EvaluationScores) -> float:
        """Compute critical value for the Friedman aligned ranks test statistic.

        Parameters
        ----------
        scores : EvaluationScores
            Evaluation scores container.

        Returns
        -------
        float
            Critical value at ``1 - alpha`` (upper-tail).
        """
        df = scores.n_algorithms - 1
        return float(chi2.ppf(1.0 - self.alpha, df))


    def _validate_scores(self, scores: EvaluationScores):
        """Validate that the input scores meet the test assumptions.

        Parameters
        ----------
        scores : EvaluationScores
            Scores to validate.

        Raises
        ------
        ValueError
            If score structure is not compatible with this test.
        """
        if scores.n_algorithms < 2:
            raise ValueError(
                "Friedman aligned ranks test requires at least two algorithms"
            )
        if scores.n_datasets < 2:
            raise ValueError(
                "Friedman aligned ranks test requires at least two datasets"
            )

    def _get_metadata(self) -> dict[str, Any]:
        """Return procedure metadata.

        Returns
        -------
        dict
            Metadata including ranking convention.
        """
        return {"minimize": self.minimize}

    def __validate_arguments(self, minimize: bool):
        """Validate initialization arguments.

        Parameters
        ----------
        minimize : bool
            Whether to minimize the metric (lower scores are better).
        """
        if not isinstance(minimize, bool):
            raise ValueError("minimize must be a boolean")