Drift diffusion models: how people trade speed for accuracy, and what they should do

def simulate_ddm(v, a, w, t0, n, rng, sv=0.0, sz=0.0, dt=1e-3, t_max=6.0, correct_bias=True):
    """Euler simulation of the DDM. Parameters may be scalars or length-n arrays.

    Drift ~ N(v, sv) per trial; start uniform on a * (w +- sz / 2). Returns response times
    (NaN if no threshold is reached by t_max) and whether the upper threshold was hit."""
    v, a, w, t0, sv, sz = (np.broadcast_to(np.asarray(p, float), (n,)) for p in (v, a, w, t0, sv, sz))
    drift = v + sv * rng.standard_normal(n)
    x = a * (w + sz * (rng.random(n) - 0.5))
    shift = 0.5826 * np.sqrt(dt) if correct_bias else 0.0
    rt, upper, alive = np.full(n, np.nan), np.zeros(n, bool), np.arange(n)
    for step in range(1, int(t_max / dt) + 1):
        x[alive] += drift[alive] * dt + np.sqrt(dt) * rng.standard_normal(alive.size)
        up, down = x[alive] >= a[alive] - shift, x[alive] <= shift
        rt[alive[up | down]] = step * dt + t0[alive[up | down]]
        upper[alive[up]] = True
        alive = alive[~(up | down)]
        if alive.size == 0:
            break
    return rt, upper


# a few paths for the picture, with the same scheme
V0, A0, W0, T00 = 1.5, 1.2, 0.5, 0.3
fig, axes = plt.subplots(3, 1, figsize=(10, 6.5), height_ratios=[1, 2.2, 1], sharex=True)
path_rng = np.random.default_rng(3)
dt = 2e-3
for i in range(24):
    x, xs = A0 * W0, [A0 * W0]
    while 0 < x < A0 and len(xs) < 2000:
        x += V0 * dt + np.sqrt(dt) * path_rng.standard_normal()
        xs.append(min(max(x, 0), A0))
    tt = T00 + dt * np.arange(len(xs))
    axes[1].plot(tt, xs, color=BLUE if xs[-1] >= A0 else ORANGE, lw=0.8, alpha=0.7)
axes[1].axhline(A0, color="k", lw=1.2)
axes[1].axhline(0, color="k", lw=1.2)
axes[1].axvspan(0, T00, color=GREY, alpha=0.15)
axes[1].text(0.02, A0 * 0.5, "non-decision\ntime $t_0$", fontsize=9, va="center")
axes[1].annotate("", xy=(0.28, A0), xytext=(0.28, 0), arrowprops=dict(arrowstyle="<->"))
axes[1].text(0.22, A0 * 0.85, "$a$", fontsize=11)
axes[1].set(ylabel="evidence", ylim=(-0.1, A0 + 0.1))
axes[1].set_title(f"Evidence accumulates with drift v = {V0} towards 'word' (a = {A0}, w = {W0})")
rt_sim, up_sim = simulate_ddm(V0, A0, W0, T00, 40_000, np.random.default_rng(1))
bins = np.linspace(0, 2.5, 80)
axes[0].hist(rt_sim[up_sim], bins=bins, color=BLUE)
axes[0].set(ylabel="count", title=f"'word' responses ({up_sim.mean():.0%})")
axes[2].hist(rt_sim[~up_sim], bins=bins, color=ORANGE)
axes[2].invert_yaxis()
axes[2].set(ylabel="count", xlabel="response time (s)", title=f"'nonword' responses ({1 - up_sim.mean():.0%})")
axes[2].set_xlim(0, 2.5);

GL_X, GL_W = np.polynomial.legendre.leggauss(3)
U_SWITCH = 0.5


def log_f01(u, w):
    """log first-passage density at the lower threshold for a zero-drift walk between 0 and 1,
    started at w, at time u. Small-time series (k = -2..2) below U_SWITCH, large-time (k = 1, 2)
    above; each branch sees a clipped u, so neither produces inf/NaN (or NaN gradients)."""
    us, ul = pt.minimum(u, U_SWITCH), pt.maximum(u, U_SWITCH)
    small_sum = w                                           # the k = 0 term, exponent factored out
    for k in (-2, -1, 1, 2):
        small_sum = small_sum + (w + 2 * k) * pt.exp(-2 * k * (w + k) / us)
    small = pt.log(small_sum) - w**2 / (2 * us) - 0.5 * np.log(2 * np.pi) - 1.5 * pt.log(us)
    s1 = pt.sin(np.pi * w)                                  # sin(2 pi w) = 2 sin cos
    large = (np.log(np.pi) - np.pi**2 * ul / 2
             + pt.log(s1 + 4 * s1 * pt.cos(np.pi * w) * pt.exp(-1.5 * np.pi**2 * ul)))
    return pt.switch(u < U_SWITCH, small, large)


def wfpt_logp(t, upper, v, a, w, sv, sz=None, gl=(GL_X, GL_W)):
    """log density of decision time t (> 0) at the upper (True) or lower threshold.

    v: drift towards the upper threshold; a: separation; w: relative start; sv: sd of the drift
    across trials (closed form); sz: width of the uniform start range, as a fraction of a
    (Gauss-Legendre nodes on a new last axis). All arguments broadcast."""
    vv = pt.switch(upper, -v, v)
    w = pt.switch(upper, 1 - w, w)
    if sz is not None:
        gx, gw = gl
        w = w[..., None] + 0.5 * pt.switch(upper, -1, 1)[..., None] * sz[..., None] * gx
        t, vv, a = t[..., None], vv[..., None], a[..., None]
        sv = pt.as_tensor_variable(sv)
        sv = sv[..., None] if sv.ndim else sv
    lp = (log_f01(t / a**2, w) - 2 * pt.log(a)
          + ((a * w * sv) ** 2 - 2 * a * vv * w - vv**2 * t) / (2 * (1 + sv**2 * t))
          - 0.5 * pt.log1p(sv**2 * t))
    if sz is not None:
        lp = pt.logsumexp(lp + np.log(gw / 2), axis=-1)
    return lp


def compile_density(gl=(GL_X, GL_W)):
    """NumPy-callable log density of signed response times y (+ = upper) given per-trial params."""
    y, v, a, w, t0, sv, sz = (pt.dvector(n) for n in ("y", "v", "a", "w", "t0", "sv", "sz"))
    dt = pt.abs(y) - t0
    lp = wfpt_logp(pt.maximum(dt, 1e-4), y > 0, v, a, w, sv, sz, gl)
    return pytensor.function([y, v, a, w, t0, sv, sz], pt.switch(dt > 1e-4, lp, -np.inf))


ddm_logpdf = compile_density()

def density(t, upper, v, a, w, sv=0.0, sz=0.0, t0=0.0, fn=ddm_logpdf):
    t = np.asarray(t, float)
    y = np.where(upper, 1.0, -1.0) * (t + t0)
    full = lambda p: np.full(t.shape, p, float)  # noqa: E731
    return np.exp(fn(y, full(v), full(a), full(w), full(t0), full(sv), full(sz)))


t_fine = np.linspace(1e-4, 12, 120_001)
cases = [(1.5, 1.2, 0.5, 0, 0), (2.5, 0.8, 0.3, 0, 0), (-1.0, 2.0, 0.7, 0, 0),
         (2.0, 1.5, 0.5, 1.0, 0.3), (0.5, 2.5, 0.4, 0.8, 0.6)]
rows = []
for v, a, w, sv, sz in cases:
    p_up = np.trapezoid(density(t_fine, True, v, a, w, sv, sz), t_fine)
    p_lo = np.trapezoid(density(t_fine, False, v, a, w, sv, sz), t_fine)
    exact = (1 - np.exp(-2 * v * a * w)) / (1 - np.exp(-2 * v * a)) if sv == sz == 0 else np.nan
    rows.append((v, a, w, sv, sz, p_up + p_lo, p_up, exact))
print("1. total mass and P(upper)")
print(pd.DataFrame(rows, columns=["v", "a", "w", "sv", "sz", "total", "P(upper)", "closed form"]).round(6))

# 2. truncation: reference with 41 terms in each series (NumPy), switching at u = 1
u_g, w_g = np.meshgrid(np.geomspace(0.005, 8, 400), np.linspace(0.02, 0.98, 49))
k = np.arange(-20, 21)[:, None, None]
ref_small = (np.sum((w_g + 2 * k) * np.exp(-(w_g + 2 * k) ** 2 / (2 * u_g)), axis=0)
             / np.sqrt(2 * np.pi * u_g**3))
k = np.arange(1, 42)[:, None, None]
ref_large = np.pi * np.sum(k * np.exp(-(k**2) * np.pi**2 * u_g / 2) * np.sin(k * np.pi * w_g), axis=0)
ref = np.where(u_g < 1, ref_small, ref_large)
U_, W_ = pt.dmatrices("u", "w")
ours = np.exp(pytensor.function([U_, W_], log_f01(U_, W_))(u_g, w_g))
print(f"2. series truncation: max relative error {np.max(np.abs(ours / ref - 1)):.1e} "
      f"(reference series differ from each other by {np.max(np.abs(ref_small / ref_large - 1)[(u_g > 0.3) & (u_g < 2)]):.0e})")

# 3. drift variability: closed form vs 60-node Gauss-Hermite over v
gh_x, gh_w = np.polynomial.hermite_e.hermegauss(60)
t_chk = np.linspace(0.02, 4, 200)
for v, a, w, sv in [(2.0, 1.2, 0.45, 1.0), (0.5, 2.0, 0.6, 2.0)]:
    for up in (True, False):
        num = sum(wi * density(t_chk, up, v + sv * xi, a, w) for xi, wi in zip(gh_x, gh_w)) / gh_w.sum()
        cf = density(t_chk, up, v, a, w, sv)
        print(f"3. sv = {sv}, upper = {up}: max |closed form - numerical| = {np.max(np.abs(cf - num)):.1e} "
              f"(peak density {cf.max():.2f})")

# 4. start variability: 3 (and 5) vs 40 Gauss-Legendre nodes
ddm_logpdf40 = compile_density(np.polynomial.legendre.leggauss(40))
ddm_logpdf5 = compile_density(np.polynomial.legendre.leggauss(5))
for sz in (0.2, 0.35, 0.5, 0.7, 0.9):
    d40 = density(t_chk, False, 1.5, 1.2, 0.5, 0.5, sz, fn=ddm_logpdf40)
    err3 = np.max(np.abs(density(t_chk, False, 1.5, 1.2, 0.5, 0.5, sz) - d40)) / d40.max()
    err5 = np.max(np.abs(density(t_chk, False, 1.5, 1.2, 0.5, 0.5, sz, fn=ddm_logpdf5) - d40)) / d40.max()
    print(f"4. sz = {sz}: max error / peak density: 3 nodes {err3:.1e}, 5 nodes {err5:.1e}")

fit_df = df.groupby(["id", "condition", "frequency"]).sample(frac=0.25, random_state=RANDOM_SEED)
hold_df = df.drop(fit_df.index)
fit_df, hold_df = fit_df.reset_index(drop=True), hold_df.reset_index(drop=True)
print(f"fit on {len(fit_df):,} trials, hold out {len(hold_df):,}")

SUBJ = np.sort(df.id.unique())
INSTR = ["accuracy", "speed"]
LO, HI = 0.18, 3.0             # the RT window kept by the authors
GUESS_HI = 0.5                 # fast guesses: uniform on [LO, GUESS_HI], either key
LOG_GUESS = -np.log(2 * (GUESS_HI - LO))


def encode(frame):
    """Integer codes and the signed response time of each trial."""
    stim = pd.Index(STIM).get_indexer(frame["frequency"])
    return dict(subj=pd.Index(SUBJ).get_indexer(frame["id"]),
                speed=(frame["condition"] == "speed").to_numpy().astype(int),
                stim=stim, sign=np.where(stim < 3, 1.0, -1.0),
                y=np.where(frame["upper"], frame["rt"], -frame["rt"]))


def ddm_mixture_logp(y, v, a, w, t0, pc, sv, sz=None):
    """Contaminated DDM log-likelihood of signed response times (contaminant probability pc)."""
    dt = pt.abs(y) - t0
    lp = wfpt_logp(pt.maximum(dt, 1e-4), y > 0, v, a, w, sv, sz)
    lp = pt.switch(dt > 1e-4, lp, -np.inf)
    guess = pt.switch(pt.abs(y) < GUESS_HI, pt.log(pc) + LOG_GUESS, -np.inf)
    return pt.logaddexp(pt.log1p(-pc) + lp, guess)


def mix_np(lp, pc, y):
    return np.logaddexp(np.log1p(-pc) + lp, np.where(np.abs(y) < GUESS_HI, np.log(pc) + LOG_GUESS, -np.inf))


def ddm_model(frame, variability=True):
    e = encode(frame)
    coords = {"subj": SUBJ, "instr": INSTR, "stim": STIM, "trial": np.arange(len(frame))}
    with pm.Model(coords=coords) as m:
        mu_la = pm.Normal("mu_la", np.log(1.2), 0.4, dims="instr")
        sd_la = pm.HalfNormal("sd_la", 0.3)
        a = pm.Deterministic("a", pt.exp(mu_la + sd_la * pm.Normal("z_la", dims=("subj", "instr"))),
                             dims=("subj", "instr"))
        mu_v = pm.Normal("mu_v", 2.0, 1.5, dims="stim")
        sd_v = pm.HalfNormal("sd_v", 1.0)
        v = pm.Deterministic("v", mu_v + sd_v * pm.Normal("z_v", dims=("subj", "stim")), dims=("subj", "stim"))
        mu_g = pm.Normal("mu_g", 0.0, 0.3)
        sd_g = pm.HalfNormal("sd_g", 0.2)
        g = pm.Deterministic("g", mu_g + sd_g * pm.Normal("z_g", dims="subj"), dims="subj")
        mu_lt = pm.Normal("mu_lt", np.log(0.35), 0.3, dims="instr")
        sd_lt = pm.HalfNormal("sd_lt", 0.3)
        t0 = pm.Deterministic("t0", pt.exp(mu_lt + sd_lt * pm.Normal("z_lt", dims=("subj", "instr"))),
                              dims=("subj", "instr"))
        pc = pm.Beta("pc", 1.0, 49.0, dims="instr")   # probability of a fast guess, prior mean 2%
        mu_w = pm.Normal("mu_w", 0.0, 0.5)
        sd_w = pm.HalfNormal("sd_w", 0.5)
        w = pm.Deterministic("w", pm.math.invlogit(mu_w + sd_w * pm.Normal("z_w", dims="subj")), dims="subj")
        if variability:
            sv = pm.HalfNormal("sv", 1.0, dims="instr")
            sz = pm.Beta("sz", 1.5, 4.0, dims="instr")   # start range as a fraction of the widest possible
        else:
            sv, sz = pt.as_tensor(0.0), None
        s, c, j = e["subj"], e["speed"], e["stim"]
        v_i = e["sign"] * v[s, j] * pt.exp(g[s] * c)
        w_i = w[s]
        sz_i = None if sz is None else sz[c] * 2 * pt.minimum(w_i, 1 - w_i)   # stay inside (0, 1)
        sv_i = sv if sz is None else sv[c]
        params = [v_i, a[s, c], w_i, t0[s, c], pc[c], sv_i] + ([] if sz_i is None else [sz_i])
        pm.CustomDist("y", *params, logp=ddm_mixture_logp, observed=e["y"], dims="trial")
    return m

T_GRID = np.linspace(LO, HI, 500)
QS = np.array([0.1, 0.3, 0.5, 0.7, 0.9])


def variability_draws(idata, n_draws, seed):
    """Drift and start variability by instruction, shape (2, n_draws); zeros for the plain model."""
    if "sv" not in idata.posterior:
        return np.zeros((2, n_draws)), np.zeros((2, n_draws))
    return tuple(az.extract(idata, var_names=[k], num_samples=n_draws, random_seed=seed)
                 .transpose("instr", "sample").to_numpy() for k in ("sv", "sz"))


def cell_predictions(idata, frame, n_draws=100, seed=1):
    """Posterior predictive P(correct) and correct/error RT quantiles per (instruction, stimulus),
    pooling participants with the trial counts of `frame`. Returns an xarray Dataset."""
    post = az.extract(idata, var_names=["a", "v", "g", "t0", "w"], num_samples=n_draws, random_seed=seed)
    sv, sz = variability_draws(idata, n_draws, seed)
    pc = az.extract(idata, var_names=["pc"], num_samples=n_draws, random_seed=seed).transpose("instr", "sample").to_numpy()
    counts = (frame.groupby(["condition", "frequency", "id"]).size()
              .reindex(pd.MultiIndex.from_product([INSTR, STIM, SUBJ])).fillna(0).to_numpy().reshape(2, 6, -1))
    S, nT = len(SUBJ), len(T_GRID)
    # every (instr, stim, subj, threshold, time) combination for one draw
    ci, ji, si, ui, ti = np.meshgrid(np.arange(2), np.arange(6), np.arange(S), [0, 1], np.arange(nT), indexing="ij")
    ci, ji, si, ui, ti = (x.ravel() for x in (ci, ji, si, ui, ti))
    sign = np.where(ji < 3, 1.0, -1.0)
    y = np.where(ui == 1, 1.0, -1.0) * T_GRID[ti]
    out = {k: np.empty((n_draws, 2, 6, *s)) for k, s in [("p_correct", ()), ("q_correct", (5,)), ("q_error", (5,))]}
    a_, v_, g_, t0_, w_ = (post[k].transpose(..., "sample").to_numpy() for k in ["a", "v", "g", "t0", "w"])
    for d in range(n_draws):
        vi = sign * v_[si, ji, d] * np.exp(g_[si, d] * ci)
        wi = w_[si, d]
        lp = ddm_logpdf(y, vi, a_[si, ci, d], wi, t0_[si, ci, d], sv[ci, d], sz[ci, d] * 2 * np.minimum(wi, 1 - wi))
        dens = np.exp(mix_np(lp, pc[ci, d], y)).reshape(2, 6, S, 2, nT)
        pooled = np.einsum("cjsut,cjs->cjut", dens, counts) / counts.sum(-1)[..., None, None]
        cdf = np.cumsum(pooled, axis=-1) * (T_GRID[1] - T_GRID[0])
        for c in range(2):
            for j in range(6):
                up_is_correct = j < 3
                corr, err = (cdf[c, j, 1], cdf[c, j, 0]) if up_is_correct else (cdf[c, j, 0], cdf[c, j, 1])
                out["p_correct"][d, c, j] = corr[-1] / (corr[-1] + err[-1])
                out["q_correct"][d, c, j] = np.interp(QS * corr[-1], corr, T_GRID)
                out["q_error"][d, c, j] = np.interp(QS * err[-1], err, T_GRID)
    dims = {"p_correct": ("draw", "instr", "stim"), "q_correct": ("draw", "instr", "stim", "q"),
            "q_error": ("draw", "instr", "stim", "q")}
    return xr.Dataset({k: (dims[k], val) for k, val in out.items()},
                      coords={"instr": INSTR, "stim": STIM, "q": QS})


def observed_summary(frame):
    rows = {}
    for (c, j), g in frame.groupby(["condition", "frequency"]):
        rows[(c, j)] = (g.correct.mean(), np.quantile(g.rt[g.correct], QS), np.quantile(g.rt[~g.correct], QS))
    return xr.Dataset({
        "p_correct": (("instr", "stim"), [[rows[c, j][0] for j in STIM] for c in INSTR]),
        "q_correct": (("instr", "stim", "q"), [[rows[c, j][1] for j in STIM] for c in INSTR]),
        "q_error": (("instr", "stim", "q"), [[rows[c, j][2] for j in STIM] for c in INSTR])},
        coords={"instr": INSTR, "stim": STIM, "q": QS})


obs_hold = observed_summary(hold_df)
pred_plain = cell_predictions(idata_plain, hold_df)
diff_obs = 1000 * (obs_hold.q_error - obs_hold.q_correct).sel(q=0.5)
diff_plain = 1000 * (pred_plain.q_error - pred_plain.q_correct).sel(q=0.5)
print("median error RT - median correct RT (ms), held-out trials")
print(pd.DataFrame({"observed": diff_obs.to_series(),
                    "plain DDM": diff_plain.median("draw").to_series()}).unstack(0).round(0))