import numpy as np, pandas as pd, scipy, sys, os
import matplotlib.pyplot as plt, seaborn as sns
%config InlineBackend.figure_format = 'retina'

# Cache directory: slow fits are saved here so mid-talk re-runs are instant.
CACHE_DIR = "cache"
os.makedirs(CACHE_DIR, exist_ok=True)

def softmax(values, beta=1.0):
    """Turn expected values into choice probabilities.

    Bigger `beta` -> more deterministic choices (pick the best option almost every time).
    Smaller `beta` -> more random choices (choice probabilities closer to 50/50).
    """
    scaled = values * beta
    return np.exp(scaled) / np.sum(np.exp(scaled))

def rw_update(current_value, reward, alpha):
    """One Rescorla-Wagner update step: nudge a value toward the reward.

    This is the one piece of math shared by rw1a1b_sim (simulating an
    agent) and rw1a1b_fit (asking how likely a real agent's choices were
    under some alpha/beta) -- same rule, used in two directions.
    """
    prediction_error = reward - current_value
    updated_value = current_value + alpha * prediction_error
    return updated_value, prediction_error

def rw1a1b_sim(params: np.ndarray, nblocks: int = 3, ntrials: int = 24,
               outcomes: np.ndarray | None = None, seed: int | None = 42):
    """Simulate a simple Rescorla-Wagner (RW) reinforcement-learning model.

    THE TASK
    --------
    On every trial, a simulated subject picks between two options, "A" and
    "B". Option A pays a reward 80% of the time; option B pays a reward 20%
    of the time. After choosing, the subject updates their belief about how
    good the *chosen* option is (its "expected value"), using the
    Rescorla-Wagner learning rule:

        new_value = old_value + alpha * (reward - old_value)

    - `alpha` (learning rate) controls how much the subject updates after
      each outcome.
    - `beta` (inverse temperature) controls how deterministically the
      subject picks the option they currently think is best (see softmax
      above).

    PARAMETERS
    ----------
    params : np.ndarray, shape (n_subjects, 2)
        Column 0 = beta  (inverse temperature; must be in [1e-5, 20])
        Column 1 = alpha (learning rate; must be in [0, 1])
    nblocks : int
        Number of separate task blocks to simulate per subject.
    ntrials : int
        Number of trials per block.
    outcomes : np.ndarray or None, shape (nblocks, ntrials, 2)
        If given, these fixed rewards (0.0/1.0) are used instead of drawing
        rewards randomly. Handy when you want every simulated subject to
        see the *same* sequence of outcomes (e.g. real experiment data, or
        a fair comparison between subjects).
    seed : int or np.random.Generator
        Makes the simulation reproducible. Pass a fixed seed (or a
        pre-built Generator) to get identical choices/rewards every run.

    RETURNS
    -------
    A dictionary of arrays (each shaped (n_subjects, nblocks, ntrials[, 2])):
        "choices"   : "A" or "B" chosen on each trial
        "choices_A" : 1.0 if A was chosen, else 0.0
        "rewards"   : reward received (0.0 or 1.0)
        "EV"        : learned value of each option (index 0=A, 1=B).
                      One extra trial long: it also stores the value
                      *after* the final update.
        "ch_prob"   : model's choice probabilities [P(A), P(B)]
        "PE"        : prediction error (reward - expected value) that
                      drove each value update
        "nll"       : negative log-likelihood of the choice actually made
                      (used later for model fitting)
        "params"    : the (beta, alpha) used for each subject
    """
    # Reward probability for each option, as [P(reward=1), P(reward=0)].
    # Row 0 = option A (80% rewarded), row 1 = option B (20% rewarded).
    REWARD_PROBS_BY_OPTION = np.array([
        [0.8, 0.2],   # option A
        [0.2, 0.8],   # option B
    ])

    # Accept either a plain integer seed or an already-built Generator, so
    # callers can share one RNG across multiple calls if they want to.
    rng = seed if isinstance(seed, np.random.Generator) else np.random.default_rng(seed)

    n_subjects = params.shape[0]
    all_beta = params[:, 0]
    all_alpha = params[:, 1]

    # --- sanity-check the parameters -------------------------------------
    if not ((all_beta >= 1e-5) & (all_beta <= 20.0)).all():
        raise ValueError("Beta values out of bounds (must be in [1e-5, 20])")
    if not ((all_alpha >= 0.0) & (all_alpha <= 1.0)).all():
        raise ValueError("Alpha values out of bounds (must be in [0, 1])")

    # --- pre-allocate storage for everything we want to record ----------
    choices   = np.empty((n_subjects, nblocks, ntrials), dtype=object)   # "A" or "B"
    choices_A = np.zeros((n_subjects, nblocks, ntrials))                 # 1.0 / 0.0
    rewards   = np.zeros((n_subjects, nblocks, ntrials))                 # 1.0 / 0.0
    ch_prob   = np.zeros((n_subjects, nblocks, ntrials, 2))              # [P(A), P(B)]
    PE        = np.zeros((n_subjects, nblocks, ntrials))                 # prediction error
    nll       = np.zeros((n_subjects, nblocks, ntrials))                 # -log P(choice made)

    # EV ("expected value") has one *extra* trial slot per block: it stores
    # the value *before* trial 0, and *after* every trial's update.
    # Both options start at 0.5 (fully uncertain), matching pyEM convention.
    EV = np.full((n_subjects, nblocks, ntrials + 1, 2), 0.5)

    # --- run the simulation ----------------------------------------------
    for subject_i in range(n_subjects):
        beta = float(all_beta[subject_i])
        alpha = float(all_alpha[subject_i])

        for block_i in range(nblocks):
            EV[subject_i, block_i, 0, :] = 0.5  # reset values at the start of each block

            for trial_i in range(ntrials):
                current_values = EV[subject_i, block_i, trial_i, :]

                # 1) Turn current values into choice probabilities.
                choice_probs = softmax(current_values, beta)
                ch_prob[subject_i, block_i, trial_i, :] = choice_probs

                # 2) Sample a choice: 0 = A, 1 = B.
                choice = rng.choice([0, 1], p=choice_probs)
                choices[subject_i, block_i, trial_i] = "A" if choice == 0 else "B"
                choices_A[subject_i, block_i, trial_i] = 1.0 if choice == 0 else 0.0

                # 3) Get the reward for the chosen option this trial.
                if outcomes is None:
                    # No fixed schedule given -> draw a reward randomly
                    # using that option's true reward probability.
                    reward_probs = REWARD_PROBS_BY_OPTION[choice]
                    reward = rng.choice([1.0, 0.0], p=reward_probs)
                else:
                    # A fixed outcome schedule was given -> just look it up
                    # instead of drawing randomly.
                    reward = float(outcomes[block_i, trial_i, choice])
                rewards[subject_i, block_i, trial_i] = reward

                # 4) Rescorla-Wagner update: compute the prediction error
                #    and nudge the chosen option's value toward the reward.
                #    The unchosen option's value is carried over unchanged.
                new_value, prediction_error = rw_update(current_values[choice], reward, alpha)
                PE[subject_i, block_i, trial_i] = prediction_error

                EV[subject_i, block_i, trial_i + 1, :] = current_values
                EV[subject_i, block_i, trial_i + 1, choice] = new_value # current_value + alpha * prediction_error

                # 5) Record how "surprised" the model is by the choice it
                #    made -- used later when fitting the model to data.
                nll[subject_i, block_i, trial_i] = -np.log(choice_probs[choice] + 1e-12)

    return {
        "params": np.array([all_beta, all_alpha]).T,
        "choices": choices,
        "rewards": rewards,
        "EV": EV,
        "ch_prob": ch_prob,
        "choices_A": choices_A,
        "PE": PE,
        "nll": nll,
    }


# Quick reproducibility check: same seed -> identical simulated data.
_a = rw1a1b_sim(np.array([[3.0, 0.4]]), nblocks=1, ntrials=20, seed=42)
_b = rw1a1b_sim(np.array([[3.0, 0.4]]), nblocks=1, ntrials=20, seed=42)
assert np.array_equal(_a["choices"], _b["choices"]) and np.array_equal(_a["rewards"], _b["rewards"])
print("Seeded simulator is reproducible:", True)

Seeded simulator is reproducible: True

from scipy.stats import truncnorm, beta as beta_dist
# simulate computer agents completing the two-armed task
nsubjects, nblocks, ntrials = 80, 1, 100
betamin, betamax = .75, 10 # inverse temperature
alphamin, alphamax = .05, .95 # learning rate

# generate distribution of parameters within range
np.random.seed(42)
beta_rv  = truncnorm((betamin-0)/1, (betamax-0)/1, loc=.2, scale=2.2).rvs(nsubjects)
a_lo, a_hi = beta_dist.cdf([alphamin, alphamax], 1.1, 1.1)
alpha_rv = beta_dist.ppf(a_lo + np.random.rand(nsubjects)*(a_hi - a_lo), 1.1, 1.1)
rl_params = np.column_stack((beta_rv, alpha_rv))

rl_param_names = ['beta', 'alpha']
sim_output = rw1a1b_sim(rl_params, nblocks=1, ntrials=ntrials)

# plot distributions of params from all agents
fig, ax = plt.subplots(1, 2, figsize=(8, 4))
for i in range(len(rl_param_names)):
    sns.histplot(rl_params[:,i], kde=True, ax=ax[i])
    ax[i].set_title(rl_param_names[i])
plt.tight_layout()
plt.show()

# Generate two subjects: SAME beta, SAME seed (same task), only alpha differs ---
ntrials = 100
beta = 2.56
low_alpha, high_alpha = 0.05, 0.8
 
# Same seed => identical reward schedule, so the ONLY difference is the learning rate.
low  = rw1a1b_sim(np.array([[beta, low_alpha]]),  ntrials=ntrials, seed=5)
high = rw1a1b_sim(np.array([[beta, high_alpha]]), ntrials=ntrials, seed=5)

# plot actions and rewards over trials for these agents
fig, axes = plt.subplots(1, 2, figsize=(12, 4), sharey=True)
for ax, out, a, ttl in [
    (axes[0], low,  low_alpha,  f'Low learning rate (α = {low_alpha})'),
    (axes[1], high, high_alpha, f'High learning rate (α = {high_alpha})'),
]:
    # Reward: 1 / 0 per trial (blue dots)
    ax.scatter(range(ntrials), out['rewards'][0, 0, :],
               label='Reward', color='blue', alpha=.3)
    ax.set_xlabel('Trials')
    ax.set_ylim(-0.1, 1.1)
    ax.set_title(f'{ttl}\nchose the 80% arm {int(out["choices_A"][0,0,:].sum())}/{ntrials} trials')
 
    # Choice on a twin axis (red dashed line): 1 = chose arm A (80%), 0 = arm B (20%)
    ax2 = ax.twinx()
    ax2.plot(range(ntrials), out['choices_A'][0, 0, :].astype(int),
             label='Choice', color='red', linestyle='--')
    ax2.set_ylim(-0.1, 1.1)
    ax2.set_yticks([0, 1])
    ax2.set_yticklabels(['Arm B (20%)', 'Arm A (80%)'])
 
axes[0].set_ylabel('Rewards (Blue)', color='blue')
# only label the right-hand choice axis on the rightmost panel to avoid clutter
axes[1].figure.axes[-1].set_ylabel('Choices (Red)', color='red', rotation=270, labelpad=28)
fig.suptitle(r'Choices & rewards — same task, same $\beta$, only $\alpha$ differs', fontsize=12)
plt.tight_layout()
plt.savefig('choices_rewards_low_vs_high.png', dpi=150, bbox_inches='tight')

# also plot Q-values over trials
fig, axes = plt.subplots(1, 2, figsize=(11, 4), sharey=True)
for ax, out, a, ttl in [
    (axes[0], low,  low_alpha,  f'Low learning rate (α = {low_alpha})'),
    (axes[1], high, high_alpha, f'High learning rate (α = {high_alpha})'),
]:
    ax.plot(range(ntrials), out['EV'][0, 0, 1:, 0], label=r'$Q_{1}$ (80% arm)', color='blue')
    ax.plot(range(ntrials), out['EV'][0, 0, 1:, 1], label=r'$Q_{2}$ (20% arm)',
            color='red', linestyle='--')
    ax.axhline(0.8, color='0.6', lw=1, ls=':', label='true reward rate (0.8)')
    ax.set_xlabel('Trials')
    ax.set_title(ttl)
    sns.despine(ax=ax)
axes[0].set_ylabel('Q-Value')
axes[0].legend(loc='center right', fontsize=8, frameon=False)
fig.suptitle(r'Same task, same $\beta$ — only $\alpha$ differs', fontsize=12)
plt.tight_layout()
plt.savefig('low_vs_high_alpha.png', dpi=150, bbox_inches='tight')
print("Low-alpha  Q1: min=%.3f max=%.3f std=%.3f" %
      (low['EV'][0,0,1:,0].min(), low['EV'][0,0,1:,0].max(), low['EV'][0,0,1:,0].std()))
print("High-alpha Q1: min=%.3f max=%.3f std=%.3f" %
      (high['EV'][0,0,1:,0].min(), high['EV'][0,0,1:,0].max(), high['EV'][0,0,1:,0].std()))

Low-alpha  Q1: min=0.500 max=0.826 std=0.093
High-alpha Q1: min=0.008 max=1.000 std=0.408

# Let's generate some data based on a simple linear model
np.random.seed(42)
n = 100
x1 = np.random.uniform(0, 10, n)
x2 = np.random.uniform(0, 10, n)
noise = np.random.normal(0, 1, n)
b0 =  2.41 # intercept
b1 =  1.64 # slope 1
b2 = -3.27 # slope 2
y = b0 + b1*x1 + b2*x2 + noise

# Let's plot the data
fig, ax = plt.subplots()
ax.scatter(x1, y, label='x1')
ax.scatter(x2, y, label='x2')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend(loc='lower center')
sns.despine()
plt.title('Simulated Data')
plt.show()

# let's create a function to calculate the negative log-likelihood - this is often called an objective function
def negll_lm(params, y, x1, x2):
    b0, b1, b2 = params
    y_pred = b0 + b1*x1 + b2*x2
    likelihoods = scipy.stats.norm.logpdf(y, y_pred)
    return -np.sum(likelihoods)

# we can use scipy to minimize the negative log-likelihood
initial_guess = [0, 0, 0]
result = scipy.optimize.minimize(negll_lm,
                                 initial_guess, 
                                 args=(y, x1, x2))
b0_hat, b1_hat, b2_hat = result.x

# grid search: try many combinations of b0, b1, b2 and keep the best (lowest negative log-likelihood)
b0_grid = np.linspace(-5, 5, 20)
b1_grid = np.linspace(-5, 5, 20)
b2_grid = np.linspace(-5, 5, 20)

best_negll = np.inf
best_params = None
for b0_ in b0_grid:
    for b1_ in b1_grid:
        for b2_ in b2_grid:
            negll = negll_lm([b0_, b1_, b2_], y, x1, x2)
            if negll < best_negll:
                best_negll = negll
                best_params = (b0_, b1_, b2_)

b0_hat_grid, b1_hat_grid, b2_hat_grid = best_params

print(f"Grid search evaluated {len(b0_grid)*len(b1_grid)*len(b2_grid)} combinations; "
      f"BFGS took only {result.nfev} evaluations.")

Grid search evaluated 8000 combinations; BFGS took only 28 evaluations.

# let's print the estimated parameters from the optimization
print(f'Parameter | Actual | Estimated (BFGS) | Estimated (Grid)')
print(f'b0        |  {b0:.2f}  |  {b0_hat:.2f}  |  {b0_hat_grid:.2f}')
print(f'b1        |  {b1:.2f}  |  {b1_hat:.2f}  |  {b1_hat_grid:.2f}')
print(f'b2        | {b2:.2f}  | {b2_hat:.2f}  | {b2_hat_grid:.2f}')

Parameter | Actual | Estimated (BFGS) | Estimated (Grid)
b0        |  2.41  |  2.32  |  2.37
b1        |  1.64  |  1.61  |  1.84
b2        | -3.27  | -3.20  | -3.42

# we can also plot y as a function of x2 (line) with the simulated data points (scatter)
fig, ax = plt.subplots()
y_hat_bf = b0_hat + b1_hat*np.mean(x1) + b2_hat*x2
y_hat_grid = b0_hat_grid + b1_hat_grid*np.mean(x1) + b2_hat_grid*x2

# plot fitted line
ax.plot(x2, y_hat_bf, label='Predicted Data (BFGS)', color='black')
ax.plot(x2, y_hat_grid, label='Predicted Data (Grid)', color='red', linestyle='--')
ax.scatter(x2, y, label='True Simulated Data', alpha=0.5, s=50)
ax.set_xlabel('x2')
ax.set_ylabel('y')
ax.legend(loc='lower center')
sns.despine()
plt.show()

# Fit function for a single subject. Here we compute the PURE negative
# log-likelihood (our "surprise" score). The prior term is added later, in the
# MAP / EM sections -- MLE uses likelihood alone.
def norm2alpha(x):
    return scipy.special.expit(np.asarray(x))          # R -> (0, 1)

def norm2beta(x, max_val=20.0):
    return max_val / (1.0 + np.exp(-np.asarray(x)))     # R -> (0, max_val)

def calc_fval(negll, params, prior=None, output='npl'):
    """Objective value. With a prior, returns -log[P(data|h) * P(h)] (MAP);
    otherwise returns the plain negative log-likelihood (MLE)."""
    if output == 'npl' and prior is not None and hasattr(prior, 'logpdf'):
        fval = -(-negll + prior.logpdf(np.asarray(params)))
        if np.isinf(fval):
            fval = 1e7  # keep gradient-based optimizers moving
        return fval
    else:
        return negll

def rw1a1b_fit(params, choices, rewards, prior=None, output="npl"):
    """Evaluate a Rescorla-Wagner model against a subject's real choices/rewards.

    This runs the *same* learning rule as rw1a1b_sim (see rw_update), but
    instead of sampling a choice, it replays the choices/rewards a subject
    actually produced and asks: "how likely was this behavior, given this
    alpha/beta?" That likelihood (as a negative log-likelihood, `nll`) is
    what an optimizer searches over to find the alpha/beta that best
    explain the subject's data.

    `params` come from the optimizer on an *unconstrained* scale, so we
    map them into valid ranges with norm2beta/norm2alpha before using them
    (keeps beta > 0 and alpha in [0, 1] without the optimizer needing to
    respect those bounds itself).
    """
    beta = float(norm2beta(params[0]))
    alpha = float(norm2alpha(params[1]))

    # Reject impossible parameter values with a large penalty instead of
    # raising an error -- the optimizer needs a *number* to keep searching.
    if not (1e-5 <= beta <= 20.0):
        return 1e7
    if not (0.0 <= alpha <= 1.0):
        return 1e7

    nblocks, ntrials = rewards.shape
    EV = np.full((nblocks, ntrials + 1, 2), 0.5)   # init Q = 0.5, same as rw1a1b_sim
    PE = np.zeros((nblocks, ntrials))
    nll = 0.0

    for block_i in range(nblocks):
        EV[block_i, 0, :] = 0.5
        for trial_i in range(ntrials):
            # This subject's *actual* choice and reward this trial --
            # nothing is sampled here, we're just replaying what happened.
            choice = 0 if choices[block_i, trial_i] == "A" else 1
            reward = rewards[block_i, trial_i]

            choice_probs = softmax(EV[block_i, trial_i, :], beta)

            new_value, prediction_error = rw_update(
                EV[block_i, trial_i, choice], reward, alpha
            )
            PE[block_i, trial_i] = prediction_error
            EV[block_i, trial_i + 1, :] = EV[block_i, trial_i, :]
            EV[block_i, trial_i + 1, choice] = new_value

            # How surprised is the model by the choice this subject actually made?
            nll += -np.log(choice_probs[choice] + 1e-12)

    if output == "all":
        choices_A = (np.asarray(choices) == "A").astype(float)
        return {
            "params": [beta, alpha], "choices": choices, "choices_A": choices_A,
            "rewards": rewards, "EV": EV, "PE": PE, "nll": nll,
        }

    return calc_fval(nll, params, prior=prior, output=output)

# use the TD-learning model to fit the simulated data with a grid search
beta_grid_norm = np.linspace(-5, -1, 50)
beta_grid = norm2beta(beta_grid_norm)
alpha_grid_norm = np.linspace(-2, 3, 50)
alpha_grid = norm2alpha(alpha_grid_norm)
nll = np.zeros((50, 50))

for i in range(50):
    for j in range(50):
            nll[i, j] = rw1a1b_fit([beta_grid_norm[i], alpha_grid_norm[j]], sim_output['choices'][0,:,:], sim_output['rewards'][0,:,:], prior=None, output="nll")

# find the optimal parameters
i_opt, j_opt = np.unravel_index(nll.argmin(), nll.shape)
beta_opt = beta_grid[i_opt]
alpha_opt = alpha_grid[j_opt]

# plot the negative log-likelihood as a heatmap
fig, ax = plt.subplots()
sns.heatmap(nll, ax=ax, cmap='Spectral')

# Label axes
ax.set_ylabel(r'$\beta$')
ax.set_xlabel(r'$\alpha$')

# Choose a manageable number of tick labels
num_ticks = 6
beta_tick_indices  = np.linspace(0, len(beta_grid) - 1, num_ticks).astype(int)
alpha_tick_indices = np.linspace(0, len(alpha_grid) - 1, num_ticks).astype(int)

# Position ticks at cell centers
ax.set_yticks(beta_tick_indices + 0.5)
ax.set_xticks(alpha_tick_indices + 0.5)

# Apply the correct value labels
ax.set_yticklabels([f"{beta_grid[i]:.2f}"  for i in beta_tick_indices])
ax.set_xticklabels([f"{alpha_grid[j]:.2f}" for j in alpha_tick_indices])

# Ensure β increases upward, not downward
ax.invert_yaxis()

# mark the optimal parameters
ax.scatter(j_opt, i_opt, color='black', marker='x', s=100)

ax.collections[0].colorbar.set_label('Negative Log-Likelihood (lower = better)', rotation=270, labelpad=20)

plt.show()

# print the estimated parameters from the grid search
print(f'Parameter | Actual | Estimated')
print(f'beta      |  {sim_output["params"][0,0]:.2f}   |  {beta_opt:.2f}')
print(f'alpha     |  {sim_output["params"][0,1]:.2f}  |  {alpha_opt:.2f}')

Parameter | Actual | Estimated
beta      |  2.56   |  2.80
alpha     |  0.82  |  0.72

# now let's use scipy to minimize the negative log-likelihood
initial_guess = [.2, 1]
result = scipy.optimize.minimize(rw1a1b_fit, initial_guess, args=(sim_output['choices'][0,:,:], sim_output['rewards'][0,:,:], None, 'nll'))
beta_hat, alpha_hat = norm2beta(result.x[0]), norm2alpha(result.x[1])

# print the estimated parameters from the optimization
print(f'Parameter | Actual | Estimated')
print(f'beta      |  {sim_output["params"][0,0]:.2f}  |  {beta_hat:.2f}')
print(f'alpha     |  {sim_output["params"][0,1]:.2f}  |  {alpha_hat:.2f}')

Parameter | Actual | Estimated
beta      |  2.56  |  2.89
alpha     |  0.82  |  0.73

# Fit a SEPARATE MLE model to each agent independently (no pooling).
# Cached: slow to re-run live, so we save/load. The seeded sim guarantees
# the cache matches a live re-fit.
_mle_cache = os.path.join(CACHE_DIR, "mle_indep_params.npy")
if os.path.exists(_mle_cache):
    rl_est_params = np.load(_mle_cache)
    print("Loaded cached independent-MLE estimates.")
else:
    rl_est_params = np.zeros((nsubjects, 2))
    rng_init = np.random.default_rng(0)
    for simS in range(nsubjects):
        initial_guess = rng_init.normal(-1, 1, 2)
        result = scipy.optimize.minimize(
            rw1a1b_fit, initial_guess,
            args=(sim_output['choices'][simS, :, :], sim_output['rewards'][simS, :, :], None, 'nll'))
        rl_est_params[simS, :] = norm2beta(result.x[0]), norm2alpha(result.x[1])
    np.save(_mle_cache, rl_est_params)
    print("Fit and cached independent-MLE estimates.")

Loaded cached independent-MLE estimates.

# Recovery scatter for independent MLE. This is panel 1 of the running
# "worse -> better -> best" comparison.
from pyem.utils import plotting

fig, axes = plt.subplots(1, 2, figsize=(8, 4))
for i, name in enumerate(rl_param_names):
    plotting.plot_scatter(rl_params[:, i], f'Simulated {name}',
                          rl_est_params[:, i], f'Estimated {name}',
                          ax=axes[i])
fig.suptitle('MLE (independent per-subject fits)')
plt.tight_layout()
plt.show()

# Where MLE breaks: with few trials, independent MLE estimates get noisy and
# pile up on the boundaries. We PRE-RENDER this figure once (seeded, offline)
# and load the PNG so the talk has no live re-simulation here.
_ss_png = os.path.join("figures", "mle_smallsample.png")
os.makedirs("figures", exist_ok=True)

if not os.path.exists(_ss_png):
    ss_ntrials = 25   # short experiment
    ss_sim = rw1a1b_sim(rl_params, nblocks=1, ntrials=ss_ntrials, seed=123)
    ss_est = np.zeros((nsubjects, 2))
    _rng = np.random.default_rng(7)
    for s in range(nsubjects):
        ig = _rng.normal(-1, 1, 2)
        r = scipy.optimize.minimize(
            rw1a1b_fit, ig,
            args=(ss_sim['choices'][s, :, :], ss_sim['rewards'][s, :, :], None, 'nll'))
        ss_est[s, :] = norm2beta(r.x[0]), norm2alpha(r.x[1])

    figss, axss = plt.subplots(1, 2, figsize=(8, 4))
    for i, name in enumerate(rl_param_names):
        plotting.plot_scatter(rl_params[:, i], f'Simulated {name}',
                              ss_est[:, i], f'Estimated {name}',
                              ax=axss[i], colorname='firebrick')
    figss.suptitle(f'MLE with only {ss_ntrials} trials/subject — recovery fails')
    figss.tight_layout()
    figss.savefig(_ss_png, dpi=150, bbox_inches='tight')
    plt.close(figss)

from IPython.display import Image, display
display(Image(filename=_ss_png))

# Standalone MAP fit: re-fit the SAME subjects, but now with a FIXED prior.
# We use pyEM's native fixed-prior support (prior_mu / prior_sigma) -- same
# EMModel, same fit function, just a prior instead of learning it from data.
from pyem import EMModel

# Build per-subject data rows [choices, rewards] from the simulated dataset.
map_data = [[sim_output['choices'][s], sim_output['rewards'][s]]
            for s in range(nsubjects)]

# A gentle prior in the UNCONSTRAINED (normalized) space:
#   beta ~ mid-range, alpha ~ 0.5. sigma=1 keeps it soft (a rubber band, not a clamp).
prior_mu    = np.array([0.0, 0.0])
prior_sigma = np.array([1.0, 1.0])

_map_cache = os.path.join(CACHE_DIR, "map_fixedprior_params.npy")
map_model = EMModel(all_data=map_data, fit_func=rw1a1b_fit,
                    param_names=["beta", "alpha"],
                    param_xform=[norm2beta, norm2alpha])
if os.path.exists(_map_cache):
    map_est_params = np.load(_map_cache)
    # still need the fitted model object for its posterior; refit is cheap enough,
    # but to stay instant we only reload the estimates for plotting.
    print("Loaded cached MAP estimates.")
else:
    # mstep_maxit=1 keeps the prior FIXED: a single MAP pass at the prior we
    # specified, rather than letting EM re-learn the group prior. This is the
    # genuine 'MAP with a fixed prior' beat (EM comes next and learns it).
    map_model.fit(prior_mu=prior_mu, prior_sigma=prior_sigma, verbose=0,
                  mstep_maxit=1, max_restarts=1)
    map_est_params = map_model.subject_params()   # natural space (beta, alpha)
    np.save(_map_cache, map_est_params)
    print("Fit and cached MAP estimates.")

Loaded cached MAP estimates.

# Overlay MAP estimates on the MLE recovery scatter. The prior pulls the wild,
# boundary-slamming MLE points inward toward plausible values.
fig, axes = plt.subplots(1, 2, figsize=(8, 4))
for i, name in enumerate(rl_param_names):
    axes[i].scatter(rl_params[:, i], rl_est_params[:, i],
                    color='0.6', alpha=0.4, s=45, label='MLE')
    axes[i].scatter(rl_params[:, i], map_est_params[:, i],
                    color='seagreen', alpha=0.6, s=45, label='MAP')
    lo = min(rl_params[:, i].min(), 0)
    hi = rl_params[:, i].max()
    axes[i].plot([lo, hi], [lo, hi], 'k--', lw=1)
    axes[i].set_xlabel(f'Simulated {name}')
    axes[i].set_ylabel(f'Estimated {name}')
    axes[i].set_title(name)
    if i == 0:
        axes[i].legend(loc='upper left', frameon=False)
    sns.despine(ax=axes[i])
fig.suptitle('MAP (fixed prior) vs MLE — boundary points pulled inward')
plt.tight_layout()
plt.show()

# Hierarchical EM fit + parameter recovery. EM learns the group prior FROM THE
# DATA (no hand-picked prior), then re-fits everyone under it. Cached because
# the 80-subject EM loop is the slowest step in the talk.
_em_true   = os.path.join(CACHE_DIR, "em_true_params.npy")
_em_est    = os.path.join(CACHE_DIR, "em_est_params.npy")
_em_mu     = os.path.join(CACHE_DIR, "em_posterior_mu.npy")

em_model = EMModel(all_data=None, fit_func=rw1a1b_fit,
                   param_names=["beta", "alpha"],
                   param_xform=[norm2beta, norm2alpha],
                   simulate_func=rw1a1b_sim)

if os.path.exists(_em_est) and os.path.exists(_em_mu):
    em_true_params = np.load(_em_true)
    em_est_params  = np.load(_em_est)
    em_posterior_mu = np.load(_em_mu)   # NORMALIZED (Gaussian) space
    print("Loaded cached EM recovery.")
else:
    recovery = em_model.recover(rl_params, pr_inputs=['choices', 'rewards'],
                                nblocks=nblocks, ntrials=ntrials, seed=42)
    em_true_params = recovery['true_params']
    em_est_params  = recovery['estimated_params']            # natural space
    em_posterior_mu = recovery['recovery_model'].posterior()['mu']  # normalized space
    np.save(_em_true, em_true_params)
    np.save(_em_est, em_est_params)
    np.save(_em_mu, em_posterior_mu)
    print("Fit and cached EM recovery.")

# Third recovery scatter (best of the three)
fig, axes = plt.subplots(1, 2, figsize=(8, 4))
for i, name in enumerate(rl_param_names):
    plotting.plot_scatter(em_true_params[:, i], f'Simulated {name}',
                          em_est_params[:, i], f'Estimated {name}',
                          ax=axes[i], colorname='seagreen')
fig.suptitle('Hierarchical EM — best recovery')
plt.tight_layout()
plt.show()

Loaded cached EM recovery.

# Draw the shrinkage: MLE (independent) vs EM (hierarchical), with the group
# center EM learned. CORRECTNESS: em_posterior_mu is in NORMALIZED space, so we
# transform it through norm2beta / norm2alpha before plotting in natural units.
group_center_natural = np.array([norm2beta(em_posterior_mu[0]),
                                 norm2alpha(em_posterior_mu[1])])
print(f"EM group center (natural): beta={group_center_natural[0]:.2f}, "
      f"alpha={group_center_natural[1]:.2f}")

# The y-axis is arbitrary here — jitter just spreads overlapping points apart so
# the shrinkage arrows are legible. Each subject gets ONE y-value, shared by its
# MLE point, its EM point, and its arrow, so arrows stay horizontal.
_jit_rng = np.random.default_rng(0)
jitter_h = 0.4  # vertical spread; purely cosmetic

fig, axes = plt.subplots(1, 2, figsize=(9, 4))
for i, name in enumerate(rl_param_names):
    y = _jit_rng.uniform(-jitter_h, jitter_h, nsubjects)
    # arrows from each MLE estimate to its EM estimate (the rubber band)
    for s in range(nsubjects):
        axes[i].annotate('', xy=(em_est_params[s, i], y[s]),
                         xytext=(rl_est_params[s, i], y[s]),
                         arrowprops=dict(arrowstyle='->', color='0.75', lw=0.6))
    axes[i].scatter(rl_est_params[:, i], y,
                    color='0.6', s=30, label='MLE', zorder=3)
    axes[i].scatter(em_est_params[:, i], y,
                    color='seagreen', s=30, label='EM', zorder=3)
    axes[i].axvline(group_center_natural[i], color='crimson', ls='--',
                    label='EM group center')
    axes[i].set_ylim(-1, 1)          # give the jitter some breathing room
    axes[i].set_yticks([])
    axes[i].set_xlabel(name)
    axes[i].set_title(f'{name}: estimates shrink toward the group center')
    if i == 0:
        axes[i].legend(loc='upper right', frameon=False, fontsize=8)
    sns.despine(ax=axes[i], left=True)
plt.tight_layout()
plt.show()

EM group center (natural): beta=2.99, alpha=0.47

# A compact linear-regression dataset + one EM fit, purely to make the
# "random effects / fixed effects / EM" analogy concrete against statsmodels.
# (This reuses the SAME EMModel machinery we used for the RL model.)
def map_lm(params, y, X, prior=None, output='npl'):
    y_pred = np.dot(X, params)
    negll = -np.sum(scipy.stats.norm.logpdf(y, y_pred))
    if output == 'nll':
        return negll
    if output == 'all':
        return {'params': params, 'y_pred': y_pred, 'negll': negll}
    if output == 'npl' and prior is not None and hasattr(prior, 'logpdf'):
        return -(-negll + prior.logpdf(np.asarray(params)))
    return negll

lm_param_names = ['b0', 'b1', 'b2']
lm_nparams = len(lm_param_names)
lm_nsubjects, lm_ntrials = 100, 100

_rng = np.random.default_rng(2021)
lm_params = _rng.normal(0, 1, size=(lm_nsubjects, lm_nparams))
X_out = np.zeros((lm_nsubjects, lm_ntrials, lm_nparams))
Y_out = np.zeros((lm_nsubjects, lm_ntrials))
for s in range(lm_nsubjects):
    X_out[s, :, 0] = 1.0
    X_out[s, :, 1:] = _rng.normal(size=(lm_ntrials, lm_nparams - 1))
    Y_out[s, :] = X_out[s] @ lm_params[s] + _rng.normal(size=lm_ntrials)
lm_all_data = [[y, x] for y, x in zip(Y_out, X_out)]

_lm_cache = os.path.join(CACHE_DIR, "lm_em_posterior_mu.npy")
if os.path.exists(_lm_cache):
    lm_posterior_mu = np.load(_lm_cache)
    print("Loaded cached LR EM group means.")
else:
    lm_model = EMModel(all_data=lm_all_data, fit_func=map_lm,
                       param_names=lm_param_names)
    lm_fit = lm_model.fit(verbose=0)
    lm_posterior_mu = lm_fit.posterior_mu
    np.save(_lm_cache, lm_posterior_mu)
    print("Fit and cached LR EM group means.")

Loaded cached LR EM group means.

# Fit a mixed-effects model to the SAME data with statsmodels, then compare its
# fixed effects to pyEM's group means (posterior_mu). They should closely agree.
from statsmodels.regression.mixed_linear_model import MixedLM

rows = []
for s in range(lm_nsubjects):
    for t in range(lm_ntrials):
        rows.append([s, X_out[s, t, 0], X_out[s, t, 1], X_out[s, t, 2], Y_out[s, t]])
df = pd.DataFrame(rows, columns=['subject', 'x0', 'x1', 'x2', 'Y'])

me_result = MixedLM.from_formula('Y ~ 1 + x1 + x2', df, groups='subject').fit()

print(f'Parameter | Mixed-effects fixed effect | pyEM posterior_mu (group mean)')
print(f'b0        |          {me_result.params["Intercept"]:+.3f}          |      {lm_posterior_mu[0]:+.3f}')
print(f'b1        |          {me_result.params["x1"]:+.3f}          |      {lm_posterior_mu[1]:+.3f}')
print(f'b2        |          {me_result.params["x2"]:+.3f}          |      {lm_posterior_mu[2]:+.3f}')
print("\nSame idea, two toolboxes: random effects = individual params, "
      "fixed effect = group mean, EM = the estimator.")

Parameter | Mixed-effects fixed effect | pyEM posterior_mu (group mean)
b0        |          -0.050          |      -0.047
b1        |          -0.048          |      -0.026
b2        |          +0.135          |      +0.149

Same idea, two toolboxes: random effects = individual params, fixed effect = group mean, EM = the estimator.

# Summarizing improvements across methods: same recovery, three ways, worse -> better -> best.
methods = [('MLE (no pooling)', rl_est_params, '0.5'),
           ('MAP (fixed prior)', map_est_params, 'seagreen'),
           ('EM (learned prior)', em_est_params, 'crimson')]

fig, axes = plt.subplots(2, 3, figsize=(12, 8))
for col, (title, est, color) in enumerate(methods):
    for row, name in enumerate(rl_param_names):
        ax = axes[row, col]
        r, _ = scipy.stats.pearsonr(rl_params[:, row], est[:, row])
        ax.scatter(rl_params[:, row], est[:, row], color=color, alpha=0.55, s=40)
        lo = min(rl_params[:, row].min(), 0); hi = rl_params[:, row].max()
        ax.plot([lo, hi], [lo, hi], 'k--', lw=1)
        ax.set_xlabel(f'Simulated {name}')
        ax.set_ylabel(f'Estimated {name}')
        ax.annotate(fr'$r={r:.2f}$', xy=(0.05, 0.9), xycoords='axes fraction')
        if row == 0:
            ax.set_title(title)
        sns.despine(ax=ax)
fig.suptitle('One model, one recovery plot, three estimators: worse → better → best',
             fontsize=13)
plt.tight_layout()
plt.show()

# Generate posterior predictive simulations from EM fits
# Let the fitted model play the task itself, many times, to build a
# DISTRIBUTION of predicted behavior. em_est_params holds the EM/MAP
# estimates in natural (beta, alpha) space — exactly what rw1a1b_sim wants.
n_ppc = 200   # number of synthetic datasets (more = smoother bands)

ppc_choicesA = np.zeros((n_ppc, nsubjects, ntrials))
for rep in range(n_ppc):
    # a fresh seed each rep => fresh choices/rewards => captures model stochasticity
    rep_sim = rw1a1b_sim(em_est_params, nblocks=nblocks, ntrials=ntrials, seed=1000 + rep)
    ppc_choicesA[rep] = rep_sim['choices_A'][:, 0, :]

# Observed behavior (the data we actually fit)
obs_choicesA = sim_output['choices_A'][:, 0, :]        # (nsubjects, ntrials)


# Learning curve: observed data vs model prediction
# Summary stat: P(choosing the optimal 80% arm) at each trial, averaged
# over subjects. The model's prediction is a BAND (2.5-97.5 percentile
# across the n_ppc synthetic datasets), not a single line.
obs_curve       = obs_choicesA.mean(axis=0)            # (ntrials,)
ppc_curve_bysim = ppc_choicesA.mean(axis=1)            # (n_ppc, ntrials)
ppc_mean = ppc_curve_bysim.mean(axis=0)
ppc_lo   = np.percentile(ppc_curve_bysim, 2.5,  axis=0)
ppc_hi   = np.percentile(ppc_curve_bysim, 97.5, axis=0)

# bin into blocks of `w` trials so the observed curve isn't jagged
w = 5
def _bin(x): return x[:(len(x)//w)*w].reshape(-1, w).mean(1)
tb = np.arange(len(_bin(obs_curve))) * w + w/2

fig, ax = plt.subplots(figsize=(6, 4))
ax.fill_between(tb, _bin(ppc_lo), _bin(ppc_hi), color='crimson', alpha=.2,
                label='95% predictive band')
ax.plot(tb, _bin(ppc_mean), color='crimson', label='Model (PPC mean)')
ax.plot(tb, _bin(obs_curve), 'o-', color='black', zorder=3, label='Observed')
ax.axhline(0.5, color='0.7', ls=':')
ax.set_ylim(0, 1)
ax.set_xlabel('Trial')
ax.set_ylabel('P(choose optimal arm)')
ax.set_title('Learning curve: does the model reproduce the data?')
ax.legend(fontsize=8, frameon=False)
sns.despine()
plt.tight_layout()
plt.show()

# how often does the observed curve land inside the model's 95% band?
coverage = np.mean((obs_curve >= ppc_lo) & (obs_curve <= ppc_hi))
print(f"Observed learning curve inside 95% predictive band on "
      f"{coverage*100:.0f}% of trials.")

# Per-subject check: predicted vs observed preference
# Summary stat per subject: overall P(choosing the optimal arm). If the
# points hug the diagonal, the model reproduces individual differences.
obs_popt = obs_choicesA.mean(axis=1)                   # (nsubjects,)
ppc_popt = ppc_choicesA.mean(axis=(0, 2))             # mean over reps & trials
ppc_popt_lo = np.percentile(ppc_choicesA.mean(axis=2), 2.5,  axis=0)
ppc_popt_hi = np.percentile(ppc_choicesA.mean(axis=2), 97.5, axis=0)

r = np.corrcoef(obs_popt, ppc_popt)[0, 1]

fig, ax = plt.subplots(figsize=(5, 5))
ax.errorbar(obs_popt, ppc_popt,
            yerr=[ppc_popt - ppc_popt_lo, ppc_popt_hi - ppc_popt],
            fmt='o', color='crimson', alpha=.5, ecolor='0.8',
            elinewidth=0.7, ms=5)
ax.plot([0, 1], [0, 1], 'k--', lw=1)
ax.set_xlim(0, 1); ax.set_ylim(0, 1)
ax.set_xlabel('Observed P(optimal), per subject')
ax.set_ylabel('Predicted P(optimal), per subject')
ax.set_title(f'Per-subject posterior predictive fit (r = {r:.2f})')
sns.despine()
plt.tight_layout()
plt.show()

Observed learning curve inside 95% predictive band on 99% of trials.

	MLE (no pooling)	MAP (fixed prior)	Hierarchical EM
Variance	high (noisy)	lower	lowest
Bias	~none	some (toward your prior)	some (toward group)
Prior	none	you pick it (arbitrary)	learned from data
Assumptions	fewest	prior chosen by hand	group distribution is Gaussian-ish
Uncertainty	per-subject only	per-subject only	group prior treated as known (understated)
Compute	cheap	cheap	iterative (moderate)

Fitting models to behavioral data¶

Who am I?¶

Goals of this tutorial¶

3 - 2 - 1 Exercise¶

Section 1 — Goals and Notes¶

Section 2 — The Rescorla–Wagner model on a two-armed bandit¶

A reinforcement-learning model (Rescorla–Wagner)¶

Two-Armed Bandit Task¶

Muddiest Point (3 minutes)¶

Section 3 — Likelihood & MLE¶

What is Maximum Likelihood Estimation (MLE)?¶

Key Concepts¶

Example 2: Linear Regression Model¶

Algorithms for Maximizing Likelihood¶

Estimating the RL parameters with MLE¶

Estimating Parameters with MLE¶

Discuss (3 minutes)¶

Section 4 — Where MLE breaks¶

Section 5 — Priors & the MAP¶

What is Maximum A Posteriori (MAP) Estimation?¶

Section 6 — Hierarchical EM¶

Hierarchical modeling with Expectation-Maximization (EM)¶

Shrinkage: what EM actually bought us¶

Section 7 — Trade-offs & a critical-reader checklist¶

Trade-offs: what shrinkage costs¶

A critical-reader checklist¶

Reflection (3 minutes)¶

Close¶

Appendix¶

Model Comparison¶

Posterior Predictive Checks¶

Additional Resources¶