draw_fig.py

from types import SimpleNamespace
import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.metrics import roc_curve, auc
from typing import List, Dict

from utils import is_correct_answer, load_json_hdf5


def plot_Z_hist(Zs:np.ndarray, plt_config:SimpleNamespace, title:str=None, filename:str='hist.png'):
    plt.figure(figsize=plt_config.figsize)
    plt.hist(Zs, bins=plt_config.bins, density=True, alpha=0.7)
    plt.xlabel("$Z_n$ (coding rate) [bit]")
    plt.ylabel("density")
    plt.xlim(plt_config.surprisal_int)
    if title is not None:
        plt.title(title)
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()

def plot_Z_ecdf(Zs:np.ndarray, plt_config:SimpleNamespace, title:str=None,filename:str='ECDF.png'):
    Zs_sorted = np.sort(Zs)
    n = len(Zs)
    y = np.arange(1, n+1) / n

    plt.figure(figsize=plt_config.figsize)
    plt.plot(Zs_sorted, y)
    plt.xlabel("$Z_n$ (coding rate) [bit]")
    plt.ylabel("ECDF")
    plt.xlim(plt_config.surprisal_int)
    if title is not None:
        plt.title(title)
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()

def plot_Z_hist_with_bounds(Zs:np.ndarray, plt_config:SimpleNamespace, 
                            title:str=None, filename:str='density.png') -> tuple[float,float]:
    p_limsup_est = np.quantile(Zs, 1-plt_config.q_thres)
    p_liminf_est = np.quantile(Zs, plt_config.q_thres)

    plt.figure(figsize=plt_config.figsize)
    plt.hist(Zs, bins=plt_config.bins, density=True, alpha=0.7)
    mean = Zs.mean()
    plt.axvline(mean, linestyle=":", label=f"mean: {mean:.3f}", color="C0")
    plt.axvline(p_limsup_est, linestyle="-.", label=fr"p-limsup/inf ~ $\theta$={plt_config.q_thres}", color="C0")
    plt.axvline(p_liminf_est, linestyle="-.", label=None, color="C0")
    if plt_config.surprisal_int is not None:
        plt.xlim(plt_config.surprisal_int)
    plt.xlabel("$Z_n$ (coding rate) [bit]")
    plt.ylabel("density")
    if title is not None:
        plt.title(title)
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()




def summarize_spectrum(Zs:np.ndarray, plt_config:SimpleNamespace, ns:List[int],
                       titles:List[str]=None, filenames:List[str]=['sum_spec.png','sum_spec2.png']):

    def estimate_stats(Zsample, q_thres):
        p_sup = np.quantile(Zsample, 1 - q_thres)
        p_inf = np.quantile(Zsample, q_thres)
        width = p_sup - p_inf
        mean = Zsample.mean()
        return p_sup, p_inf, width, mean

    # 設定
    B = 200  # サンプリング回数（100〜500くらいでOK）
    rng = np.random.default_rng(0)

    # 結果格納（平均）
    p_sup_mean, p_inf_mean, width_mean, mean_mean = [], [], [], []
    # 5-95% CI
    p_sup_lo, p_sup_hi = [], []
    p_inf_lo, p_inf_hi = [], []
    width_lo, width_hi = [], []
    mean_lo, mean_hi = [], []

    for n in ns:
        # 各 n で B回ブートストラップ
        stats = np.zeros((B, 4))
        for b in range(B):
            # with replacement
            idx = rng.integers(0, len(Zs), size=n)
            Zb = Zs[idx]
            stats[b] = estimate_stats(Zb, plt_config.q_thres)

        # stats[:,0]=p_sup, [:,1]=p_inf, [:,2]=width, [:,3]=mean

        # 平均（推定値の代表）
        p_sup_mean.append(stats[:,0].mean())
        p_inf_mean.append(stats[:,1].mean())
        width_mean.append(stats[:,2].mean())
        mean_mean.append(stats[:,3].mean())

        # 90% CI（5-95%）
        p_sup_lo.append(np.quantile(stats[:,0], 0.05))
        p_sup_hi.append(np.quantile(stats[:,0], 0.95))

        p_inf_lo.append(np.quantile(stats[:,1], 0.05))
        p_inf_hi.append(np.quantile(stats[:,1], 0.95))

        width_lo.append(np.quantile(stats[:,2], 0.05))
        width_hi.append(np.quantile(stats[:,2], 0.95))

        mean_lo.append(np.quantile(stats[:,3], 0.05))
        mean_hi.append(np.quantile(stats[:,3], 0.95))

    if filenames[0] is not None:
        # 幅の収束
        plt.figure(figsize=plt_config.figsize)
        plt.plot(ns, width_mean, marker="o")
        plt.fill_between(ns, width_lo, width_hi, alpha=0.2)

        plt.xlabel("$n$")
        plt.ylabel("width = p-limsup - p-liminf")
        if titles is not None:
            plt.title(titles[0])
        plt.grid(True)
        plt.tight_layout()
        plt.savefig(filenames[0])
        plt.close()

    if filenames[1] is not None:
        # 上・下・平均の推移
        plt.figure(figsize=plt_config.figsize)
        # p-limsup
        plt.plot(ns, p_sup_mean, marker="o", label="p-limsup")
        plt.fill_between(ns, p_sup_lo, p_sup_hi, alpha=0.2)        
        # p-liminf
        plt.plot(ns, p_inf_mean, marker="o", label="p-liminf")
        plt.fill_between(ns, p_inf_lo, p_inf_hi, alpha=0.2)
        # mean
        plt.plot(ns, mean_mean, marker="o", label="mean $Z_n$")
        plt.fill_between(ns, mean_lo, mean_hi, alpha=0.2)
        plt.xlabel("$n$")
        plt.ylabel("value [bit]")
        if plt_config.surprisal_int is not None:
            plt.ylim(plt_config.surprisal_int)
        if titles is not None:
            plt.title(titles[1])
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
        plt.savefig(filenames[1])
        plt.close()


def plot_Z_hist_correct_incorrect(Z_correct:np.ndarray, Z_incorrect:np.ndarray, plt_config:SimpleNamespace,
                                  title:str=None, filename:str='hist_correct_incorrect.png'):
    plt.figure(figsize=plt_config.figsize)

    if len(Z_correct) > 0:
        plt.hist(Z_correct, bins=plt_config.bins, alpha=0.6, density=True, label="correct", color="C0")
    if len(Z_incorrect) > 0:
        plt.hist(Z_incorrect, bins=plt_config.bins, alpha=0.6, density=True, label="incorrect", color="C1")

    # それぞれの平均や分位点も線で表示
    for Zs, lab, col in [
        (Z_correct, "correct", "C0"),
        (Z_incorrect, "incorrect", "C1"),
    ]:
        if len(Zs) == 0:
            continue
        mean = Zs.mean()
        plt.axvline(mean, linestyle=":", label=f"{lab} mean: {mean:.3f}", color=col)
        plt.axvline(np.quantile(Zs, 1-plt_config.q_thres), linestyle="-.", label=fr"{lab} p-limsup/inf ~ $\theta$={plt_config.q_thres}", color=col)
        plt.axvline(np.quantile(Zs, plt_config.q_thres), linestyle="-.", label=None, color=col)

    plt.xlabel("$Z_n$ (coding rate) [bit]")
    plt.ylabel("density")
    if plt_config.surprisal_int is not None:
        plt.xlim(plt_config.surprisal_int)
    if title is not None:
      plt.title(title)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()

def ecdf(Zs):
    Zs = np.sort(Zs)
    n = len(Zs)
    y = np.arange(1, n+1) / n
    return Zs, y

def plot_ecdf_correct_incorrect(Z_correct:np.ndarray, Z_incorrect:np.ndarray, plt_config:SimpleNamespace,
                                title:str=None,filename:str='ECDF_correct_incorrect.png'):
        
    plt.figure(figsize=plt_config.figsize)

    if len(Z_correct) > 0:
        x_c, y_c = ecdf(Z_correct)
        plt.plot(x_c, y_c, label="correct")
    if len(Z_incorrect) > 0:
        x_i, y_i = ecdf(Z_incorrect)
        plt.plot(x_i, y_i, label="incorrect")

    if plt_config.surprisal_int is not None:
        plt.xlim(plt_config.surprisal_int)
    plt.xlabel("$Z_n$ (coding rate) [bit]")
    plt.ylabel("ECDF")
    if title is not None:
        plt.title(title)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()

def plot_ecdf_correct_incorrect2(Z:dict, plt_config:SimpleNamespace,
                                title:str=None,filename:str='ECDF_correct_incorrect2.png'):
        
    plt.figure(figsize=plt_config.figsize)

    for label, color, linestyle in [
        ('correct_correct','C0','-'),
        ('incorrect_correct','C0',':'),
        ('correct_incorrect','C1',':'),
        ('incorrect_incorrect','C1','-')
        ]:
        if len(Z[label]) > 0:
            x, y = ecdf(Z[label])
            plt.plot(x, y, label=f"{label} ({len(Z[label])})", color=color, linestyle=linestyle)

    if plt_config.surprisal_int is not None:
        plt.xlim(plt_config.surprisal_int)
    plt.xlabel("$Z_n$ (coding rate) [bit]")
    plt.ylabel("ECDF")
    if title is not None:
        plt.title(title)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()


def plot_roc_curve(Zs_correct:np.ndarray, Zs_incorrect:np.ndarray, plt_config:SimpleNamespace,
                   title:str="ROC curve for error detection using surprisal (Z_n)", filename:str='roc_curve.png') -> float:
    # 正解ラベル（1が正解, 0が不正解）
    y_true = [1]*len(Zs_correct) + [0]*len(Zs_incorrect)

    # surprisal（Z_n）をスコアとして使用
    # 注意：ROC では「高いほど positive と判断される」ので、
    # 本来「不正解ほど Z_n が高い」場合は Z_n をそのまま使う
    y_score = np.concatenate([Zs_correct, Zs_incorrect])

    # ==== ROC curve の計算 ====
    fpr, tpr, thresholds = roc_curve(y_true, y_score, pos_label=0)
    # pos_label=0 にする理由：
    #  不正解ほど surprisal が高い → 不正解を「陽性」として検出したい

    roc_auc = auc(fpr, tpr)

    # ==== プロット ====
    plt.figure(figsize=plt_config.figsize)
    plt.plot(fpr, tpr, label=f"ROC curve (AUC = {roc_auc:.3f})", linewidth=2)
    plt.plot([0,1], [0,1], "k--", label="random baseline")

    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    if title is not None:
        plt.title(title)
    plt.legend()
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()

    # ==== optional: しきい値を確認 ====
    # threshold と TPR/FPR を対応づけて見たい場合：
    roc_table = pd.DataFrame({
        "threshold": thresholds,
        "TPR": tpr,
        "FPR": fpr
    })
    #print(roc_table.head())

    return roc_auc


def main(model_name:str, temperature:float=0.6):
    exp_dir0 = f'exp_t{temperature}'

    model_name2 = os.path.basename(model_name)
    exp_dir = os.path.join(exp_dir0,model_name2)

    resdir = f'{exp_dir}/figs'
    os.makedirs(resdir,exist_ok=True)

    print(f"### {model_name} ###")
    result = {}

    plt_config = SimpleNamespace(
        figsize=(8,5), surprisal_int=None,
        q_thres=0.01, bins=60,
        )

    cond = {
        'common_QA_cond': ['common_QA','common_choices_QA'],
        'knowledge_QA_cond': ['knowledge_QA','knowledge_RAG_QA'],
        'hop_QA_cond': ['hop_QA','hop_deriv_QA'],
    }

    for gen_type in [
        'news','poem',
        'common_QA','knowledge_QA','hop_QA',
        'common_choices_QA','knowledge_RAG_QA','hop_deriv_QA',
        'common_QA_cond','knowledge_QA_cond','hop_QA_cond',
        ]:
        if not os.path.isfile(f'{exp_dir}/{gen_type}.json.gz'):
            print(f'Skip {gen_type}')
            continue
        
        print(f" {gen_type}")
        sample_log = load_json_hdf5(f'{exp_dir}/{gen_type}',read_meta_only=True)
        base_log = None
        ref_log = None
        Zn_or_dif = 'Z_n'
        plt_config.surprisal_int = (0, 3)
        Zs = {'all':[], 'correct':[], 'incorrect':[]}

        if gen_type in cond.keys():
            base_log = load_json_hdf5(f'{exp_dir}/{cond[gen_type][0]}',read_meta_only=True)
            ref_log = load_json_hdf5(f'{exp_dir}/{cond[gen_type][1]}',read_meta_only=True)
            Zn_or_dif = 'Z_n difference'
            plt_config.surprisal_int = (-3, 1)
            Zs['correct_correct'] = []
            Zs['correct_incorrect'] = []
            Zs['incorrect_correct'] = []
            Zs['incorrect_incorrect'] = []
        
        answer_not_found = 0
        for k, v in sample_log.items():
            if np.isnan(v['Z_n']):
                answer_not_found += 1
                continue
            # Zs(平均NLL)の底をe(nat)から2(bit)に変換する
            Z_n = v['Z_n'] / np.log(2)

            if ref_log is not None:
                Z_n_ref = ref_log[k]['Z_n']
                Z_n_ref = Z_n_ref / np.log(2)
                Z_n = Z_n_ref - Z_n

            Zs['all'].append(Z_n)
            if 'QA' in gen_type:
                if is_correct_answer(v):
                    Zs['correct'].append(Z_n)
                    if ref_log is not None:
                        if is_correct_answer(base_log[k]):
                            Zs['correct_correct'].append(Z_n)
                        else:
                            Zs['incorrect_correct'].append(Z_n)
                else:
                    Zs['incorrect'].append(Z_n)
                    if ref_log is not None:
                        if is_correct_answer(base_log[k]):
                            Zs['correct_incorrect'].append(Z_n)
                        else:
                            Zs['incorrect_incorrect'].append(Z_n)
            
        if answer_not_found > 0:
            print(f"{answer_not_found} answers not found")

        Zs["all"] = np.array(Zs["all"])
        n = len(Zs["all"])

        p_limsup_est = np.quantile(Zs["all"], 1-plt_config.q_thres)
        p_liminf_est = np.quantile(Zs["all"], plt_config.q_thres)

        result[gen_type] = {
            "p-limsup": p_limsup_est,
            "p-liminf": p_liminf_est,
            "mean": Zs["all"].mean(),
            "width": p_limsup_est - p_liminf_est,
        }

        if 'QA' in gen_type:
            Zs_correct = np.array(Zs["correct"])
            Zs_incorrect = np.array(Zs["incorrect"])
            plot_Z_hist_correct_incorrect(Zs_correct, Zs_incorrect, plt_config,
                            title=f"{Zn_or_dif} with estimated p-limsup/inf (n={n})",
                            filename=f'{resdir}/density_{gen_type}_correct_incorrect.png')
            
            if ref_log is None:
                plot_ecdf_correct_incorrect(Zs_correct, Zs_incorrect, plt_config,
                                title=f"ECDF of {Zn_or_dif} (n={n})",
                                filename=f'{resdir}/ECDF_{gen_type}_correct_incorrect.png')
            else:
                plot_ecdf_correct_incorrect2(Zs, plt_config,
                            title=f"ECDF of {Zn_or_dif} (n={n})",
                            filename=f'{resdir}/ECDF_{gen_type}_correct_incorrect.png')


            roc_auc = plot_roc_curve(Zs_correct, Zs_incorrect, plt_config,
                            title=f"ROC curve for error detection using surprisal ({Zn_or_dif}) (n={n})",
                            filename=f'{resdir}/roc_curve_{gen_type}.png')
            
            result[gen_type]['correct_rate'] = len(Zs_correct)/len(Zs["all"])
            result[gen_type]['roc_auc'] = roc_auc
        else:
            plot_Z_hist_with_bounds(Zs["all"], plt_config,
                            title=f"{Zn_or_dif} with estimated p-limsup/inf (n={n})",
                            filename=f'{resdir}/density_{gen_type}.png')

            #plot_Z_hist(Zs["all"], plt_config,
            #            title=f"Histogram of {Zn_or_dif} (n={n})",
            #            filename=f'{resdir}/hist_{gen_type}.png')
            
            #plot_Z_ecdf(Zs["all"], plt_config, 
            #            title=f"ECDF of {Zn_or_dif} (n={n})",
            #            filename=f'{resdir}/ECDF_{gen_type}.png')

        ns = [int(r*n) for r in np.arange(0.05,0.96,0.05)] + [n-1]

        summarize_spectrum(Zs["all"],plt_config,ns=ns,
                        titles=["Spectrum width vs n","p-limsup / p-liminf / mean vs n"],
                        filenames=[None,f'{resdir}/sum_spec2_{gen_type}.png'])
    
    pd.DataFrame(result).to_csv(exp_dir+'.csv',float_format='%.3f')


if __name__ == '__main__':

    models = [
        "tokyotech-llm/Llama-3.1-Swallow-8B-v0.5",
        #"tokyotech-llm/Llama-3.1-Swallow-8B-Instruct-v0.5",
        #"Qwen/Qwen3-8B",
        #"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
        #"deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
        #"elyza/Llama-3-ELYZA-JP-8B",
        #"mistralai/Ministral-3-8B-Base-2512",
        ]
    for model_name in models:
        for temperature in [0.6]:
        #for temperature in [0.4,0.8,1.2]:
            main(model_name, temperature=temperature)