NexQuant/rdagent/app/benchmark/factor/analysis.py

import json
import pickle
from pathlib import Path

import fire
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

from rdagent.components.benchmark.conf import BenchmarkSettings
from rdagent.components.benchmark.eval_method import FactorImplementEval


class BenchmarkAnalyzer:
    def __init__(self, settings):
        self.settings = settings
        self.index_map = self.load_index_map()

    def load_index_map(self):
        index_map = {}
        with open(self.settings.bench_data_path, "r") as file:
            factor_dict = json.load(file)
        for factor_name, data in factor_dict.items():
            index_map[factor_name] = (factor_name, data["Category"], data["Difficulty"])
        return index_map

    def load_data(self, file_path):
        file_path = Path(file_path)
        if not (file_path.is_file() and file_path.suffix == ".pkl"):
            raise ValueError("Invalid file path")

        with file_path.open("rb") as f:
            res = pickle.load(f)

        return res

    def process_results(self, results):
        final_res = {}
        for experiment, path in results.items():
            data = self.load_data(path)
            summarized_data = FactorImplementEval.summarize_res(data)
            processed_data = self.analyze_data(summarized_data)
            final_res[experiment] = processed_data.iloc[-1, :]
        return final_res

    def reformat_index(self, display_df):
        """
        reform the results from

        .. code-block:: python

                              success rate
            High_Beta_Factor           0.2

        to

        .. code-block:: python

                                                    success rate
            Category Difficulty Factor
            量价       Hard       High_Beta_Factor           0.2

        """
        new_idx = []
        display_df = display_df[display_df.index.isin(self.index_map.keys())]
        for idx in display_df.index:
            new_idx.append(self.index_map[idx])

        display_df.index = pd.MultiIndex.from_tuples(
            new_idx,
            names=["Factor", "Category", "Difficulty"],
        )
        display_df = display_df.swaplevel(0, 2).swaplevel(0, 1).sort_index(axis=0)

        return display_df.sort_index(
            key=lambda x: [{"Easy": 0, "Medium": 1, "Hard": 2, "New Discovery": 3}.get(i, i) for i in x]
        )

    def result_all_key_order(self, x):
        order_v = []
        for i in x:
            order_v.append(
                {
                    "avg. Run successful rate": 0,
                    "avg. Format successful rate": 1,
                    "avg. Correlation (value only)": 2,
                    "max. Correlation": 3,
                    "max. accuracy": 4,
                    "avg. accuracy": 5,
                }.get(i, i),
            )
        return order_v

    def analyze_data(self, sum_df):
        index = [
            "FactorSingleColumnEvaluator",
            "FactorRowCountEvaluator",
            "FactorIndexEvaluator",
            "FactorEqualValueRatioEvaluator",
            "FactorCorrelationEvaluator",
            "run factor error",
        ]
        sum_df = sum_df.reindex(index, axis=0)
        sum_df_clean = sum_df.T.groupby(level=0).apply(lambda x: x.reset_index(drop=True))

        run_error = sum_df_clean["run factor error"].unstack().T.fillna(False).astype(bool)
        succ_rate = ~run_error
        succ_rate = succ_rate.mean(axis=0).to_frame("success rate")

        succ_rate_f = self.reformat_index(succ_rate)

        # if it rasis Error when running the evaluator, we will get NaN
        # Running failures are reguarded to zero score.
        format_issue = sum_df_clean[["FactorRowCountEvaluator", "FactorIndexEvaluator"]].apply(
            lambda x: np.mean(x.fillna(0.0)), axis=1
        )
        format_succ_rate = format_issue.unstack().T.mean(axis=0).to_frame("success rate")
        format_succ_rate_f = self.reformat_index(format_succ_rate)

        corr = sum_df_clean["FactorCorrelationEvaluator"].fillna(0.0)
        corr = corr.unstack().T.mean(axis=0).to_frame("corr(only success)")
        corr_res = self.reformat_index(corr)
        corr_max = sum_df_clean["FactorCorrelationEvaluator"]

        corr_max = corr_max.unstack().T.max(axis=0).to_frame("corr(only success)")
        corr_max_res = self.reformat_index(corr_max)

        value_max = sum_df_clean["FactorEqualValueRatioEvaluator"]
        value_max = value_max.unstack().T.max(axis=0).to_frame("max_value")
        value_max_res = self.reformat_index(value_max)

        value_avg = (
            (sum_df_clean["FactorEqualValueRatioEvaluator"] * format_issue)
            .unstack()
            .T.mean(axis=0)
            .to_frame("avg_value")
        )
        value_avg_res = self.reformat_index(value_avg)

        result_all = pd.concat(
            {
                "avg. Correlation (value only)": corr_res.iloc[:, 0],
                "avg. Format successful rate": format_succ_rate_f.iloc[:, 0],
                "avg. Run successful rate": succ_rate_f.iloc[:, 0],
                "max. Correlation": corr_max_res.iloc[:, 0],
                "max. accuracy": value_max_res.iloc[:, 0],
                "avg. accuracy": value_avg_res.iloc[:, 0],
            },
            axis=1,
        )

        df = result_all.sort_index(axis=1, key=self.result_all_key_order)
        print(df)

        # Calculate the mean of each column
        mean_values = df.fillna(0.0).mean()
        mean_df = pd.DataFrame(mean_values).T

        # Assign the MultiIndex to the DataFrame
        mean_df.index = pd.MultiIndex.from_tuples([("-", "-", "Average")], names=["Factor", "Category", "Difficulty"])

        # Append the mean values to the end of the dataframe
        df_w_mean = pd.concat([df, mean_df]).astype("float")

        return df_w_mean


class Plotter:
    @staticmethod
    def change_fs(font_size):
        plt.rc("font", size=font_size)
        plt.rc("axes", titlesize=font_size)
        plt.rc("axes", labelsize=font_size)
        plt.rc("xtick", labelsize=font_size)
        plt.rc("ytick", labelsize=font_size)
        plt.rc("legend", fontsize=font_size)
        plt.rc("figure", titlesize=font_size)

    @staticmethod
    def plot_data(data, file_name, title):
        plt.figure(figsize=(10, 6))
        sns.barplot(x="index", y="b", hue="a", data=data)
        plt.xlabel("Method")
        plt.ylabel("Value")
        plt.title(title)
        plt.savefig(file_name)


def main(
    path="git_ignore_folder/eval_results/res_promptV220240724-060037.pkl",
    round=1,
    title="Comparison of Different Methods",
):
    settings = BenchmarkSettings()
    benchmark = BenchmarkAnalyzer(settings)
    results = {
        f"{round} round experiment": path,
    }
    final_results = benchmark.process_results(results)
    final_results_df = pd.DataFrame(final_results)

    Plotter.change_fs(20)
    plot_data = final_results_df.drop(["max. accuracy", "avg. accuracy"], axis=0).T
    plot_data = plot_data.reset_index().melt("index", var_name="a", value_name="b")
    Plotter.plot_data(plot_data, "./comparison_plot.png", title)


if __name__ == "__main__":
    fire.Fire(main)