analysis_results.py

from dataclasses import asdict, dataclass, field
from typing import List

import pandas as pd


@dataclass
class AnalysisPlanResults:
    """
    A dataclass used to represent the results of the experiment analysis.

    Returned by :meth:`AnalysisPlan.analyze <cluster_experiments.inference.analysis_plan.AnalysisPlan.analyze>`.
    Each row corresponds to one hypothesis test (metric × variant × dimension slice).
    Only aggregate fields are stored (ATE, p-value, etc.); for the full underlying
    model fit (e.g. statsmodels regression table) for a single run, use
    :meth:`ExperimentAnalysis.get_inference_results
    <cluster_experiments.experiment_analysis.ExperimentAnalysis.get_inference_results>`
    and then :meth:`InferenceResults.model_summary
    <cluster_experiments.experiment_analysis.InferenceResults.model_summary>`.

    Attributes
    ----------
    metric_alias : List[str]
        The alias of the metric used in the test
    control_variant_name : List[str]
        The name of the control variant
    treatment_variant_name : List[str]
        The name of the treatment variant
    control_variant_mean : List[float]
        The mean value of the control variant
    treatment_variant_mean : List[float]
        The mean value of the treatment variant
    analysis_type : List[str]
        The type of analysis performed
    ate : List[float]
        The average treatment effect
    ate_ci_lower : List[float]
        The lower bound of the confidence interval for the ATE
    ate_ci_upper : List[float]
        The upper bound of the confidence interval for the ATE
    p_value : List[float]
        The p-value of the test
    std_error : List[float]
        The standard error of the test
    dimension_name : List[str]
        The name of the dimension
    dimension_value : List[str]
        The value of the dimension
    alpha: List[float]
        The significance level of the test
    """

    metric_alias: List[str] = field(default_factory=lambda: [])
    control_variant_name: List[str] = field(default_factory=lambda: [])
    treatment_variant_name: List[str] = field(default_factory=lambda: [])
    control_variant_mean: List[float] = field(default_factory=lambda: [])
    treatment_variant_mean: List[float] = field(default_factory=lambda: [])
    analysis_type: List[str] = field(default_factory=lambda: [])
    ate: List[float] = field(default_factory=lambda: [])
    ate_ci_lower: List[float] = field(default_factory=lambda: [])
    ate_ci_upper: List[float] = field(default_factory=lambda: [])
    p_value: List[float] = field(default_factory=lambda: [])
    std_error: List[float] = field(default_factory=lambda: [])
    dimension_name: List[str] = field(default_factory=lambda: [])
    dimension_value: List[str] = field(default_factory=lambda: [])
    split_name: List[str] = field(default_factory=lambda: [])
    split_value: List[str] = field(default_factory=lambda: [])
    alpha: List[float] = field(default_factory=lambda: [])

    def __add__(self, other):
        if not isinstance(other, AnalysisPlanResults):
            return NotImplemented

        return AnalysisPlanResults(
            metric_alias=self.metric_alias + other.metric_alias,
            control_variant_name=self.control_variant_name + other.control_variant_name,
            treatment_variant_name=self.treatment_variant_name
            + other.treatment_variant_name,
            control_variant_mean=self.control_variant_mean + other.control_variant_mean,
            treatment_variant_mean=self.treatment_variant_mean
            + other.treatment_variant_mean,
            analysis_type=self.analysis_type + other.analysis_type,
            ate=self.ate + other.ate,
            ate_ci_lower=self.ate_ci_lower + other.ate_ci_lower,
            ate_ci_upper=self.ate_ci_upper + other.ate_ci_upper,
            p_value=self.p_value + other.p_value,
            std_error=self.std_error + other.std_error,
            dimension_name=self.dimension_name + other.dimension_name,
            dimension_value=self.dimension_value + other.dimension_value,
            split_name=self.split_name + other.split_name,
            split_value=self.split_value + other.split_value,
            alpha=self.alpha + other.alpha,
        )

    def to_dataframe(self, drop_empty: bool = False):
        data_dict = asdict(self)
        max_len = max(len(v) for v in data_dict.values()) if data_dict else 0
        for k, v in data_dict.items():
            if len(v) < max_len:
                data_dict[k] = v + [""] * (max_len - len(v))

        df = pd.DataFrame(data_dict)

        if drop_empty:
            defaults = {
                "dimension_name": "__total_dimension",
                "dimension_value": "total",
                "split_name": "__total_split",
                "split_value": "total",
            }
            for col, val in defaults.items():
                if col in df.columns and (df[col] == val).all():
                    df = df.drop(columns=[col])
        return df

    def __str__(self) -> str:
        n = len(self.ate)
        if n == 0:
            return "AnalysisPlanResults(0 tests)"
        ate_range = f"[{min(self.ate):.4g}, {max(self.ate):.4g}]"
        return f"AnalysisPlanResults({n} test(s): metrics={self.metric_alias}, ATE in {ate_range})"

    def summary(self) -> str:
        """Return a summary of the analysis plan results."""
        df = self.to_dataframe()
        if df.empty:
            return "Analysis plan results: no tests."
        lines = [
            "Analysis plan results",
            f"  Number of tests: {len(df)}",
            "",
            df.to_string(index=False),
        ]
        return "\n".join(lines)