explore_exploit_plot.py

import os
import glob
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from scipy.io import loadmat


BASE_DIR = os.path.dirname(os.path.abspath(__file__))   # this script's own folder

data_path     = os.path.join(BASE_DIR, "parameters", "data_with_BestColor.mat")
PAPER_FIG_DIR = os.path.join(BASE_DIR, "figures")

output_dir = os.environ.get(
    "SIM_RESULTS_DIR",
    "/Users/mingzeli/Codes/selectiveAttention/manuscript/pipeline/results",
) + os.sep

N_RUNS = 20

if not os.path.exists(data_path):
    raise SystemExit(f"Behavioural data not found:\n  {data_path}\n"
                     f"Expected a 'parameters' folder next to {os.path.basename(__file__)}.")
_n_npz = len(glob.glob(os.path.join(output_dir, "simulation_seed_*.npz")))
if _n_npz == 0:
    raise SystemExit(f"No simulation_seed_*.npz found in:\n  {output_dir}\n"
                     f"Panel B would come out blank. Set SIM_RESULTS_DIR.")
print(f"Using {_n_npz} simulation files from {output_dir}")

# ============================================================
# Style Configuration — matched to combined learning curve figure
# ============================================================
plt.rcParams.update({
    'font.family':        'Arial',
    'font.size':          11,
    'axes.titlesize':     14,
    'axes.titleweight':   'bold',
    'axes.labelsize':     12,
    'axes.labelweight':   'bold',
    'axes.linewidth':     1.5,
    'axes.spines.top':    False,
    'axes.spines.right':  False,
    'xtick.labelsize':    10,
    'ytick.labelsize':    10,
    'xtick.major.width':  1.2,
    'ytick.major.width':  1.2,
    'xtick.major.size':   5,
    'ytick.major.size':   5,
    'legend.fontsize':    9,
    'legend.frameon':     False,
    'figure.dpi':         300,
    'savefig.dpi':        300,
    'savefig.bbox':       'tight',
    'savefig.transparent': True,
})

# ============================================================
# Shared palette — matched to other figures
# ============================================================
MODEL_COLORS = {
    'None':      '#ded77a',
    'Linear':    '#a7618f',
    'Absolute':  '#8ccec2',
    'Quadratic': '#c05d56',
    'Switch':    '#e9913f',
}
RPE_ORDER    = ['None', 'Linear', 'Absolute', 'Quadratic', 'Switch']
RPE_FILE_MAP = {'None': None, 'Linear': 'linear', 'Absolute': 'abs', 'Quadratic': 'quad', 'Switch': 'switch'}

SINGLE_ALPHA = 0.9
MULTI_ALPHA  = 0.45

# Reordered to process Monkey S before Monkey B
MONKEY_STYLE = {
    'Monkey S': {'slice': slice(69, 171), 'color': 'gray',        'lw': 1.8, 'zorder': 4},
    'Monkey B': {'slice': slice(0, 69),   'color': 'saddlebrown', 'lw': 1.8, 'zorder': 5},
}

# ============================================================
# Load Data
# ============================================================
mat_data      = loadmat(data_path)['data']
options_color = mat_data['options_color'][0, 0]
target_color  = mat_data['BestColor'][0, 0].flatten()
monkey_choices = mat_data['Chosen_ind'][0, 0].flatten().astype(int) - 1
X_data        = mat_data['X_data'][0, 0]
N_trials      = target_color.size

starts = np.where(X_data[:, 4] == 1)[0]
ends   = np.append(starts[1:] - 1, N_trials - 1)

# Build previous-block target array & valid block indices
prev_block_target = np.full(N_trials, np.nan)
valid_starts, valid_ends = [], []

for mname, minfo in MONKEY_STYLE.items():
    m_starts = starts[minfo['slice']]
    m_ends   = ends[minfo['slice']]
    for i in range(1, len(m_starts)):
        prev_block_target[m_starts[i]:m_ends[i] + 1] = target_color[m_starts[i - 1]]
    valid_starts.extend(m_starts[1:])
    valid_ends.extend(m_ends[1:])

valid_starts = np.array(valid_starts)
valid_ends   = np.array(valid_ends)
L_min        = (valid_ends - valid_starts + 1).min()
x_axis       = np.arange(1, L_min + 1)

# ============================================================
# Helper: accuracy on previous-block target
# ============================================================
def get_prev_target_accuracy(choices, target_starts):
    diff  = np.abs(options_color - prev_block_target)
    dists = np.minimum(diff, 2 * np.pi - diff)
    correctness = (choices == np.argmin(dists, axis=0)).astype(float)
    correctness[np.isnan(prev_block_target)] = np.nan
    blocks = np.array([correctness[s:s + L_min] for s in target_starts])
    mean = np.nanmean(blocks, axis=0)
    sem  = np.nanstd(blocks, axis=0) / np.sqrt(blocks.shape[0])
    return mean, sem

# ============================================================
# Compute monkey curves
# ============================================================
monkey_curves = {}
for mname, minfo in MONKEY_STYLE.items():
    m_mean, m_sem = get_prev_target_accuracy(monkey_choices, starts[minfo['slice']][1:])
    monkey_curves[mname] = {'mean': m_mean, 'sem': m_sem}

# ============================================================
# Compute model curves per group × RPE type
# ============================================================
model_ee = {'Single': {}, 'Multi': {}}

for rpe_label in RPE_ORDER:
    rpe_file = RPE_FILE_MAP[rpe_label]
    for is_single in [True, False]:
        group = 'Single' if is_single else 'Multi'
        run_curves = []

        for seed in range(N_RUNS):
            fname = f"simulation_seed_{seed}_single_attention_{is_single}_rpe_type_{rpe_file}.npz"
            fpath = os.path.join(output_dir, fname)
            if not os.path.exists(fpath):
                continue
            with np.load(fpath) as md:
                m_choices = md.get('model_choice_indices', md.get('np_model_choice_indices'))
            curve_mean, _ = get_prev_target_accuracy(m_choices, valid_starts)
            run_curves.append(curve_mean)

        if run_curves:
            stack = np.vstack(run_curves)
            model_ee[group][rpe_label] = {
                'mean': stack.mean(axis=0),
                'sem':  stack.std(axis=0) / np.sqrt(len(run_curves)) if len(run_curves) > 1 else np.zeros(stack.shape[1]),
                'run_curves': run_curves,
            }

# ============================================================
# Compute per-monkey MSE (first 40 trials)
# ============================================================
n_trials_mse = 40
mse_results = {
    mname: {'single_mean': [], 'single_sem': [], 'multi_mean': [], 'multi_sem': []}
    for mname in MONKEY_STYLE
}

for rpe_label in RPE_ORDER:
    for is_single in [True, False]:
        group = 'Single' if is_single else 'Multi'
        tag   = 'single' if is_single else 'multi'

        if rpe_label not in model_ee[group]:
            for mname in MONKEY_STYLE:
                mse_results[mname][f'{tag}_mean'].append(np.nan)
                mse_results[mname][f'{tag}_sem'].append(0.0)
            continue

        run_curves = model_ee[group][rpe_label]['run_curves']

        for mname in MONKEY_STYLE:
            monkey_mean = monkey_curves[mname]['mean']
            run_mses = [np.mean((rc[:n_trials_mse] - monkey_mean[:n_trials_mse]) ** 2) for rc in run_curves]
            mse_results[mname][f'{tag}_mean'].append(np.mean(run_mses))
            mse_results[mname][f'{tag}_sem'].append(
                np.std(run_mses) / np.sqrt(len(run_mses)) if len(run_mses) > 1 else 0.0
            )

# ============================================================
# Build combined figure (2 rows × 2 cols)
# ============================================================
fig, axes = plt.subplots(
    2, 2,
    figsize=(12.6, 8),   # widened to hold the out-of-axes Panel-B legends
    gridspec_kw={'height_ratios': [1.3, 1], 'hspace': 0.35, 'wspace': 0.15},
)

ax_line_single, ax_line_multi = axes[0]
ax_bar_S,       ax_bar_B      = axes[1]  # Monkey S is now the left panel

# Shared y-axis for bottom row
ax_bar_B.sharey(ax_bar_S)

# ------------------------------------------------------------------
# Row 1 — Panel A: Explore-exploit learning curves
# ------------------------------------------------------------------
group_map = {ax_line_single: 'Single', ax_line_multi: 'Multi'}
titles_a  = {
    ax_line_single: 'Single-Focus\nAttention Models',
    ax_line_multi:  'Multi-Focus\nAttention Models',
}

for ax in [ax_line_single, ax_line_multi]:
    group = group_map[ax]

    # Monkey data (lower zorder)
    for mname, minfo in MONKEY_STYLE.items():
        mc = monkey_curves[mname]
        ax.plot(x_axis, mc['mean'], color=minfo['color'],
                linewidth=minfo['lw'], zorder=minfo['zorder'], label=mname)
        ax.fill_between(x_axis, mc['mean'] - mc['sem'], mc['mean'] + mc['sem'],
                        color=minfo['color'], alpha=0.08, zorder=minfo['zorder'] - 1)

    # Model curves (higher zorder)
    for rpe_label in RPE_ORDER:
        if rpe_label not in model_ee[group]:
            continue
        mc    = model_ee[group][rpe_label]
        color = MODEL_COLORS[rpe_label]
        disp  = f"{rpe_label} (no RPE effect)" if rpe_label == 'None' else rpe_label
        ax.plot(x_axis, mc['mean'], color=color, linewidth=1.8, label=disp, zorder=10)
        ax.fill_between(x_axis, mc['mean'] - mc['sem'], mc['mean'] + mc['sem'],
                        color=color, alpha=0.12, zorder=9)

    ax.set_ylim(0, 1.0)
    ax.set_xlim(0, L_min + 1)
    ax.set_xlabel('Trial Number within Block', fontsize=12, fontweight='bold')
    ax.set_title(titles_a[ax], fontsize=14, fontweight='bold')
    ax.tick_params(axis='both', which='major', length=5, width=1.2)

ax_line_single.set_ylabel('Prev-Target Accuracy', fontsize=12, fontweight='bold')
ax_line_multi.set_yticklabels([])

# Legend on right panel — two-part: monkeys (patches) + RPE color patches
legend_monkey = [
    Patch(facecolor='gray',        edgecolor='black', linewidth=0.8, label='Monkey S'),
    Patch(facecolor='saddlebrown', edgecolor='black', linewidth=0.8, label='Monkey B'),
]
leg1_a = ax_line_multi.legend(handles=legend_monkey, loc='upper right', fontsize=8,
                               frameon=False, handlelength=1.2, handleheight=0.9,
                               labelspacing=0.3)

legend_rpe_patches = [
    Patch(facecolor=MODEL_COLORS['None'],      edgecolor='black', linewidth=0.8, label='None'),
    Patch(facecolor=MODEL_COLORS['Linear'],    edgecolor='black', linewidth=0.8, label='Linear'),
    Patch(facecolor=MODEL_COLORS['Absolute'],  edgecolor='black', linewidth=0.8, label='Absolute'),
    Patch(facecolor=MODEL_COLORS['Quadratic'], edgecolor='black', linewidth=0.8, label='Quadratic'),
    Patch(facecolor=MODEL_COLORS['Switch'],    edgecolor='black', linewidth=0.8, label='Switch'),
]
ax_line_multi.legend(handles=legend_rpe_patches, loc='upper right', fontsize=8,
                      frameon=False, handlelength=1.2, handleheight=0.9,
                      labelspacing=0.3,
                      bbox_to_anchor=(0.62, 1.0))
ax_line_multi.add_artist(leg1_a)

# ------------------------------------------------------------------
# Row 2 — Panel B: MSE bar plots (per monkey)
# ------------------------------------------------------------------
x_bar     = np.arange(len(RPE_ORDER))
bar_width = 0.35

for ax, mname in zip([ax_bar_S, ax_bar_B], ['Monkey S', 'Monkey B']):
    r = mse_results[mname]

    # Single-focus bars (solid)
    for i, rpe_label in enumerate(RPE_ORDER):
        color = MODEL_COLORS[rpe_label]
        ax.bar(x_bar[i] - bar_width / 2, r['single_mean'][i], bar_width,
               yerr=r['single_sem'][i],
               color=color, alpha=SINGLE_ALPHA,
               edgecolor='black', linewidth=0.8,
               capsize=3, error_kw={'linewidth': 1.0, 'capthick': 1.0})

    # Multi-focus bars (lighter + hatched)
    for i, rpe_label in enumerate(RPE_ORDER):
        color = MODEL_COLORS[rpe_label]
        ax.bar(x_bar[i] + bar_width / 2, r['multi_mean'][i], bar_width,
               yerr=r['multi_sem'][i],
               color=color, alpha=MULTI_ALPHA,
               edgecolor='black', linewidth=0.8, hatch='///',
               capsize=3, error_kw={'linewidth': 1.0, 'capthick': 1.0})

    ax.set_xticks(x_bar)
    ax.set_xticklabels(RPE_ORDER, fontsize=9)
    ax.set_xlabel('RPE-Attention Type', fontsize=11, fontweight='bold')
    ax.set_title(mname, fontsize=14, fontweight='bold')
    ax.tick_params(axis='both', which='major', length=5, width=1.2)

ax_bar_S.set_ylabel(f'Explore-Exploit MSE\n(First {n_trials_mse} Trials)', fontsize=11, fontweight='bold')
ax_bar_B.tick_params(axis='y', labelleft=False)

# Bar legend — two-part: focus group + RPE colors (Added to the rightmost plot, ax_bar_B)
legend_focus = [
    Patch(facecolor='gray', alpha=SINGLE_ALPHA, edgecolor='black',
          linewidth=0.8, label='Single-Focus'),
    Patch(facecolor='gray', alpha=MULTI_ALPHA, edgecolor='black',
          linewidth=0.8, hatch='///', label='Multi-Focus'),
]

legend_rpe_bar = [
    Patch(facecolor=MODEL_COLORS['None'],      edgecolor='black', linewidth=0.8, label='None'),
    Patch(facecolor=MODEL_COLORS['Linear'],    edgecolor='black', linewidth=0.8, label='Linear'),
    Patch(facecolor=MODEL_COLORS['Absolute'],  edgecolor='black', linewidth=0.8, label='Absolute'),
    Patch(facecolor=MODEL_COLORS['Quadratic'], edgecolor='black', linewidth=0.8, label='Quadratic'),
    Patch(facecolor=MODEL_COLORS['Switch'],    edgecolor='black', linewidth=0.8, label='Switch'),
]

# (Panel-B legends are created after tight_layout - see below.)

# ------------------------------------------------------------------
# Title and panel labels (A, B)
# ------------------------------------------------------------------
fig.suptitle('Explore-Exploit Dynamics Analysis', fontsize=16, fontweight='bold', y=1.01)
fig.text(0.01, 0.97, 'A', fontsize=18, fontweight='bold', va='top')
fig.text(0.01, 0.46, 'B', fontsize=18, fontweight='bold', va='top')

plt.tight_layout(rect=[0, 0, 0.86, 0.97])


_pos = ax_bar_B.get_position()
_lx  = _pos.x1 + 0.018                      # just clear of the axes' right edge
fig.legend(handles=legend_rpe_bar, loc='upper left', fontsize=8,
           frameon=False, handlelength=1.2, handleheight=0.9, labelspacing=0.3,
           bbox_to_anchor=(_lx, _pos.y1))
fig.legend(handles=legend_focus, loc='upper left', fontsize=8,
           frameon=False, handlelength=1.2, handleheight=0.9, labelspacing=0.3,
           bbox_to_anchor=(_lx, _pos.y0 + 0.42 * _pos.height))

# ------------------------------------------------------------------
# Save
# ------------------------------------------------------------------
plots_dir = PAPER_FIG_DIR
os.makedirs(plots_dir, exist_ok=True)

pdf_path = os.path.join(plots_dir, 'explore_exploit_plot.pdf')
png_path = os.path.join(plots_dir, 'explore_exploit_plot.png')
plt.savefig(pdf_path, format='pdf')
plt.savefig(png_path, format='png', dpi=300)
print(f"Saved: {pdf_path}")
print(f"Saved: {png_path}")
plt.close()