plot_stats.py

import os
import pickle
import yaml
import matplotlib.pyplot as plt
import numpy as np
from tueplots import bundles, figsizes
from argparse import ArgumentParser
from utils import create_result_path

plt.rcParams.update(bundles.iclr2024())

EXP_CONFIGS = {
    1: {
        "TS_rand wRS": "exp1.1",
        "TS_hypo wRS": "exp1.2_2.2_3.3_4.1_5.2_6.3_7.4",
    },
    2: {
        "Naive Student": "exp2.1",
        "Rational Student": "exp1.2_2.2_3.3_4.1_5.2_6.3_7.4",
    },
    3: {
        "SS_random": "exp3.1",
        "SS_uncertainty": "exp3.2",
        "SS_hypothesis": "exp1.2_2.2_3.3_4.1_5.2_6.3_7.4",
        "Lazy Student": "exp3.4",
    },
    4: {
        "TA_naive": "exp4.4",
        "TA_random": "exp4.3",
        "TA_uncertainty": "exp4.2",
        "TA_hypothesis": "exp1.2_2.2_3.3_4.1_5.2_6.3_7.4",
    },
    5: {
        r"$\alpha, \beta = 0.1$": "exp5.1",
        r"$\alpha, \beta = 1$": "exp1.2_2.2_3.3_4.1_5.2_6.3_7.4",
        r"$\alpha, \beta = 10$": "exp5.3",
    },
    6: {
        r"$K = 10$": "exp6.1",
        r"$K = 50$": "exp6.2",
        r"$K = 100$": "exp1.2_2.2_3.3_4.1_5.2_6.3_7.4",
    },
    7: {
        "LT_SS_rand": "exp7.1",
        "LT_SS_uncert": "exp7.2",
        "LT_SS_hypo": "exp7.3",
        "Active Interaction": "exp1.2_2.2_3.3_4.1_5.2_6.3_7.4",
    },
}


class ExpConfigurations:
    seed: int
    n_hypotheses: int
    n_clusters: int
    n_features: int
    n_samples: int
    data_initialization: str
    interaction_mode: str
    teacher_strategy: str
    teacher_n_beliefs: int
    teacher_alpha: float
    teacher_student_mode_assumption: str
    teacher_student_strategy_assumption: str = ""
    student_beta: float
    student_mode: str
    student_strategy: str = ""
    student_teacher_strategy_assumption: str = ""
    result_dir: str


def load_config(config_file, args):
    config_path = os.path.join("configs", f"{config_file}.yaml")
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    exp_config = ExpConfigurations()
    for key, value in config.items():
        setattr(exp_config, key, value)
    exp_config.result_dir = args.result_dir
    return exp_config


def main(args):
    result_probs = {}
    result_ranks = {}
    for exp_name, config_file in EXP_CONFIGS[args.exp].items():
        if args.env == "medium":
            config_file += "_m"
        elif args.env == "difficult":
            config_file += "_d"
        print(f"Processing {exp_name} from {config_file}")
        exp_config = load_config(config_file, args)
        for seed in args.seeds:
            exp_config.seed = seed
            result_file = create_result_path(exp_config)
            with open(result_file, "rb") as f:
                result_buffer = pickle.load(f)

            # Extract true hypothesis probabilities
            student_true_hypothesis_probs = np.array(
                result_buffer["student_true_hypothesis_probs"]
            )  # (n_steps,)
            student_true_hypothesis_probs[0] = 1 / len(result_buffer["hypotheses"])

            if exp_name not in result_probs:
                result_probs[exp_name] = []
            result_probs[exp_name].append(
                student_true_hypothesis_probs[: args.n_rounds + 1]
            )

            # Calculate iterations to reach rank #1
            iterations_to_rank1 = np.where(
                np.array(result_buffer["student_true_hypothesis_ranks"]) == 1
            )[0]
            if len(iterations_to_rank1) == 0:
                iterations_to_rank1 = len(student_true_hypothesis_probs)
            else:
                iterations_to_rank1 = (
                    iterations_to_rank1[0] + 1
                )  # +1 to convert index to iteration count

            # iterations_to_rank1 = np.where(student_true_hypothesis_probs >= 0.95)[0]
            # if len(iterations_to_rank1) == 0:
            #     iterations_to_rank1 = len(student_true_hypothesis_probs)
            # else:
            #     iterations_to_rank1 = (
            #         iterations_to_rank1[0] + 1
            #     )  # +1 to convert index to iteration count

            if exp_name not in result_ranks:
                result_ranks[exp_name] = []
            result_ranks[exp_name].append(iterations_to_rank1)

    # Draw 1 plot containing 2 subplots
    # Line plot: Probability of true hypothesis over time
    # Column plot: Iteration to achieve rank #1 (with mean and std)
    fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2)
    iclr_size = figsizes.iclr2024()["figure.figsize"]
    fig.set_size_inches(iclr_size[0] * 1.2, iclr_size[1] * 1.2)
    color_cycle = plt.rcParams["axes.prop_cycle"].by_key()["color"]

    # Line plot
    for i, (exp_name, probs) in enumerate(result_probs.items()):
        probs = np.array(probs, dtype=float)  # (n_seeds, n_steps)
        mean_probs = np.mean(probs, axis=0)
        std_probs = np.std(probs, axis=0, ddof=1)
        std_error_probs = std_probs * 1.96 / np.sqrt(probs.shape[0])  # 95% CI
        # Plot mean with markers every 10 steps
        color = color_cycle[i % len(color_cycle)]
        ax1.plot(mean_probs, label=exp_name, marker="o", markevery=10, color=color)
        ax1.fill_between(
            range(len(mean_probs)),
            mean_probs - std_error_probs,
            mean_probs + std_error_probs,
            alpha=0.2,
        )
    ax1.set_xlabel("Iteration")
    ax1.set_ylabel("Probability")
    ax1.set_title("Probability of True Hypothesis in Student Belief")
    ax1.legend()
    ax1.set_ylim(0, 1.1)
    ax1.axhline(y=0.95, color="r", linestyle="--")
    ax1.text(10.5, 0.96, r"95\% Probability", color="r")

    # Column plot
    exp_names = list(result_ranks.keys())
    iterations = [result_ranks[exp_name] for exp_name in exp_names]
    mean_iterations = [np.mean(iters) for iters in iterations]
    std_iterations = [np.std(iters, ddof=1) for iters in iterations]
    std_error_iterations = [
        std * 1.96 / np.sqrt(len(iters))
        for std, iters in zip(std_iterations, iterations)
    ]
    bar_colors = [color_cycle[i % len(color_cycle)] for i in range(len(exp_names))]
    ax2.bar(
        exp_names,
        mean_iterations,
        yerr=std_error_iterations,
        capsize=5,
        color=bar_colors,
    )
    ax2.set_ylabel(r"Iteration")
    ax2.set_title(r"True Hypothesis Reaches Rank \#1 in Student Belief")
    # ax2.set_title(r"True Hypothesis Reaches 95\% in Student Belief")

    plt.savefig(
        os.path.join(args.result_dir, f"exp_{args.exp}_{args.env}.png"), dpi=300
    )
    plt.close()


if __name__ == "__main__":
    parser = ArgumentParser(description="Teaching Simulation")
    parser.add_argument(
        "--exp", type=int, choices=[1, 2, 3, 4, 5, 6, 7], help="Experiment type"
    )
    parser.add_argument(
        "--env",
        type=str,
        choices=["easy", "medium", "difficult"],
        default="easy",
        help="Environment type",
    )
    parser.add_argument(
        "--seeds",
        type=int,
        nargs="+",
        default=[
            2,
            3,
            5,
            7,
            11,
            13,
            17,
            19,
            23,
            29,
            31,
            37,
            41,
            43,
            47,
            53,
            59,
            61,
            67,
            71,
        ],
        help="Random seeds",
    )
    parser.add_argument(
        "--n_rounds",
        type=int,
        default=100,
        help="Number of rounds for each simulation",
    )
    parser.add_argument(
        "--result_dir",
        type=str,
        default="results",
        help="Output directory to save simulation results",
    )
    args = parser.parse_args()

    os.makedirs(args.result_dir, exist_ok=True)
    main(args)