plot_trials.py

import argparse
import csv
import json
import logging
import os
from pathlib import Path

import matplotlib.pyplot as plt

from alno.utils.loading import setup_logging
from alno.utils.plotting import (
    plot_averaged_cumulative_regret,
    plot_cumulative_max,
    plot_cumulative_regret,
    plot_summary_across_trials,
)
from run_trials import EXPERIMENT_SCRIPTS


def load_results(results_dirs, algorithms, n_trials):
    """
    Load outcome and regret data from all specified directories and trials.
    Returns:
        algo_outcomes: dict of algo -> list of list of outcomes per trial
        algo_regret: dict of algo -> list of list of regrets per trial
    """
    algo_outcomes = {algo.lower(): [] for algo in algorithms}
    algo_regret = {algo.lower(): [] for algo in algorithms}

    for results_dir in results_dirs:
        for trial in range(1, n_trials + 1):
            trial_dir = Path(results_dir) / f"trial_{trial}"
            for algo in algorithms:
                algo = algo.lower()
                csv_path = trial_dir / f"pool_{algo}.csv"
                if csv_path.exists():
                    outcomes = []
                    regret = []
                    with open(csv_path, "r") as f:
                        reader = csv.DictReader(f)
                        for row in reader:
                            try:
                                outcomes.append(float(row["outcome"]))
                                regret.append(float(row["regret"]))
                            except Exception:
                                logger.warning(
                                    f"Skipping invalid row in {csv_path}: {row}"
                                )
                                continue
                    if outcomes:
                        # Append to algo list; for multiple directories, last overwrite is naturally handled
                        if len(algo_outcomes[algo]) < trial:
                            algo_outcomes[algo].append(outcomes)
                            algo_regret[algo].append(regret)
                        else:
                            # Overwrite for current trial
                            algo_outcomes[algo][trial - 1] = outcomes
                            algo_regret[algo][trial - 1] = regret
                    else:
                        logger.warning(f"No valid outcomes in {csv_path}")
                else:
                    logger.warning(f"CSV does not exist: {csv_path}")
    return algo_outcomes, algo_regret


def main(
    results_dirs,
    algorithms,
    n_trials,
    max_iterations,
    figdpi,
    n_std,
    truncate_to_shortest,
):
    # Save the directories used into a JSON file
    # Load data
    logger.info("Loading results from directories: %s", results_dirs)
    algo_outcomes, algo_regret = load_results(results_dirs, algorithms, n_trials)

    # Create output directory if it doesn't exist
    out_dir = args.output_dir
    if out_dir is None:
        out_dir = results_dirs[0]
    os.makedirs(out_dir, exist_ok=True)

    # Save the directories used into a JSON file
    json_path = os.path.join(out_dir, "plot_trials_data_dirs.json")
    with open(json_path, "w") as json_file:
        json.dump({"results_dirs": results_dirs, "n_std": n_std}, json_file)
    logger.info(f"Saved list of data directories to {json_path}")

    # Plotting
    colors = plt.get_cmap("tab10").colors
    figsize = (5, 3)

    plt.figure(figsize=figsize)
    plot_summary_across_trials(algo_outcomes, colors, n_std=n_std)
    plt.legend()
    plt.title("Average outcomes across trials")
    plt.xlim(0, max_iterations)
    output_path = os.path.join(out_dir, "summary_across_trials.png")
    plt.savefig(output_path, dpi=figdpi)
    logger.info("Saved plot: %s", output_path)
    plt.close()

    plt.figure(figsize=figsize)
    plot_cumulative_max(algo_outcomes, colors, n_std=n_std)
    plt.legend()
    plt.title("Cumulative max (best-so-far) across trials")
    plt.xlim(0, max_iterations)
    output_path = os.path.join(out_dir, "cumulative_max_across_trials.png")
    plt.savefig(output_path, dpi=figdpi)
    logger.info("Saved plot: %s", output_path)
    plt.close()

    plt.figure(figsize=figsize)
    plot_cumulative_regret(
        algo_regret, colors, n_std=n_std, truncate_to_shortest=truncate_to_shortest
    )
    plt.legend()
    plt.title("Cumulative regret across trials")
    plt.xlim(0, max_iterations)
    output_path = os.path.join(out_dir, "cumulative_regret_across_trials.png")
    plt.savefig(output_path, dpi=figdpi)
    logger.info("Saved plot: %s", output_path)
    plt.close()

    plt.figure(figsize=figsize)
    regrets_summary = plot_averaged_cumulative_regret(
        algo_regret, colors, n_std=n_std, truncate_to_shortest=truncate_to_shortest
    )
    plt.legend()
    plt.title("Averaged cumulative regret across trials")
    plt.xlim(0, max_iterations)
    output_path = os.path.join(out_dir, "averaged_cumulative_regret.png")
    plt.savefig(output_path, dpi=figdpi)
    logger.info("Saved plot: %s", output_path)
    plt.close()

    # Save regrets summary to JSON
    summary_json_path = os.path.join(out_dir, "average_regrets.json")
    with open(summary_json_path, "w") as json_file:
        json.dump(regrets_summary, json_file, indent=4)
    logger.info(f"Saved regrets summary to {summary_json_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate summary plots from run_trials results"
    )
    parser.add_argument(
        "results_dirs",
        nargs="+",
        type=str,
        help="Paths to one or more results directories",
    )
    parser.add_argument(
        "--algorithms", nargs="+", type=str, default=None, help="List of algorithms"
    )
    parser.add_argument("--n-trials", type=int, default=10, help="Number of trials")
    parser.add_argument(
        "--max-iterations", type=int, default=None, help="Maximum number of iterations"
    )
    parser.add_argument("--figdpi", type=int, default=300, help="DPI for the plot")
    parser.add_argument(
        "--n-std",
        type=int,
        default=1,
        help="Number of standard deviations for confidence intervals",
    )
    parser.add_argument(
        "--truncate-to-shortest",
        action="store_true",
        help="Truncate data to shortest trial, instead of padding",
    )
    parser.add_argument(
        "-o", "--output-dir", type=str, default=None, help="Output directory"
    )
    args = parser.parse_args()

    setup_logging()
    logger = logging.getLogger(__name__)

    # If algorithms not provided, infer from CSV filenames or set default
    if args.algorithms is None:
        algorithms = list(EXPERIMENT_SCRIPTS.keys())
        logger.info(f"Inferred algorithms: {algorithms}")
    else:
        algorithms = args.algorithms
        logger.info(f"Using specified algorithms: {algorithms}")

    try:
        main(
            args.results_dirs,
            algorithms,
            args.n_trials,
            args.max_iterations,
            args.figdpi,
            args.n_std,
            args.truncate_to_shortest,
        )
    except Exception:
        logger.exception("Unexpected error in main function")