Bayesian optimisation: choosing the next experiment

def to_grid(v):
    """Flat vector over the 1,728 reactions -> 36 x 48 image: rows ligand x temperature,
    columns solvent x base x concentration."""
    return np.asarray(v).reshape(NL).transpose(0, 4, 2, 1, 3).reshape(36, 48)


def grid_pos(i):
    """Row and column of reaction i in the to_grid image."""
    l, b, s, c, t = np.unravel_index(i, NL)
    return l * 3 + t, s * 12 + b * 3 + c


def draw_space(ax, v, cmap="viridis", vmin=None, vmax=None, labels=True, title=None, cbar=None):
    im = ax.imshow(to_grid(v), cmap=cmap, vmin=vmin, vmax=vmax, aspect="auto",
                   interpolation="nearest")
    for r in range(3, 36, 3):
        ax.axhline(r - 0.5, color="w", lw=0.6)
    for c in range(3, 48, 3):
        ax.axvline(c - 0.5, color="w", lw=0.3 if c % 12 else 1.6)
    ax.set_xticks([]), ax.set_yticks([])
    if labels:
        ax.set_yticks(np.arange(1, 36, 3), LEV[0], fontsize=7)
        ax.set_xticks(np.arange(5.5, 48, 12), LEV[2], fontsize=8)
        ax.xaxis.tick_top()
    if title:
        ax.set_xlabel(title, fontsize=9)
    if cbar is not None:
        plt.colorbar(im, ax=ax, shrink=0.8, label=cbar)
    return im


fig, ax = plt.subplots(figsize=(11, 7))
draw_space(ax, Y, vmin=0, vmax=100, cbar="yield (%)",
           title="each ligand band: 90 / 105 / 120 C (top to bottom); each solvent block: bases "
                 + " / ".join(LEV[1]) + ", each at 3 concentrations (low to high)")
best = np.argsort(Y)[-10:]
rr, cc = grid_pos(best)
ax.plot(cc, rr, "s", mfc="none", mec="red", ms=7, mew=1.2)
ax.set_title("The whole design space: all 1,728 yields (red squares: the 10 reactions above 95%)",
             fontsize=10, pad=28);

POS = [None, None, None,
       (np.array(LEV[3]) - LEV[3][0]) / (LEV[3][-1] - LEV[3][0]),   # concentration scaled to 0..1
       (np.array(LEV[4]) - LEV[4][0]) / (LEV[4][-1] - LEV[4][0])]   # temperature scaled to 0..1


def sq_dist(encoding="categorical", lig_desc=None):
    """Squared 'distances' between the levels of each factor (the kernel is rho ** d2)."""
    out = []
    for f in range(5):
        if f == 0 and encoding == "number":        # ligand as its position in an alphabetical list
            p = np.arange(NL[0]) / (NL[0] - 1)
            out.append((p[:, None] - p[None, :]) ** 2)
        elif f == 0 and encoding == "descriptors":
            out.append(lig_desc)
        elif POS[f] is None:                       # categorical: every pair of levels is 'distance 1'
            out.append(1.0 - np.eye(NL[f]))
        else:
            out.append((POS[f][:, None] - POS[f][None, :]) ** 2)
    return out


def gp_model(kernel="additive", encoding="categorical", lig_desc=None, n0=5):
    d2 = sq_dist(encoding, lig_desc)
    with pm.Model(coords={"factor": FACT}) as m:
        idx = pm.Data("idx", IDX[:n0])
        y = pm.Data("y", np.zeros(n0))
        c = pm.Normal("c", 0.0, 1.0)
        sigma = pm.HalfNormal("sigma", 0.5)
        rho = pm.Beta("rho", 2.0, 2.0, dims="factor")
        eta = pm.HalfNormal("eta", 1.5)
        K = eta**2                                  # interactions: a product over the factors
        for f in range(5):
            R = rho[f] ** pt.as_tensor(d2[f])
            K = K * R[idx[:, f][:, None], idx[:, f][None, :]]
        if kernel == "additive":                    # plus one main-effect kernel per factor
            amp = pm.HalfNormal("amp", 1.0, dims="factor")
            for f in range(5):
                Rm = pt.as_tensor(np.eye(NL[f])) if np.array_equal(d2[f], 1.0 - np.eye(NL[f])) \
                    else rho[f] ** pt.as_tensor(d2[f])
                K = K + amp[f] ** 2 * Rm[idx[:, f][:, None], idx[:, f][None, :]]
        cov = K + (sigma**2 + 1e-6) * pt.eye(idx.shape[0])
        pm.MvNormal("obs", mu=c, cov=cov, observed=y)
    return m, d2

def level_mats(h, d2):
    """Per-factor level correlation matrices: interaction (R) and main-effect (M) kernels."""
    R = [h["rho"][f] ** d2[f] for f in range(5)]
    M = [np.eye(NL[f]) if np.array_equal(d2[f], 1.0 - np.eye(NL[f])) else R[f] for f in range(5)]
    return R, M


def kern(h, d2, A, B):
    R, M = level_mats(h, d2)
    K = np.full((len(A), len(B)), h["eta"] ** 2)
    for f in range(5):
        K *= R[f][A[:, f][:, None], B[:, f][None, :]]
    if "amp" in h:
        for f in range(5):
            K += h["amp"][f] ** 2 * M[f][A[:, f][:, None], B[:, f][None, :]]
    return K


def prior_var(h):
    return h["eta"] ** 2 + (np.sum(h["amp"] ** 2) if "amp" in h else 0.0)


def gp_posterior(h, d2, sel, ys):
    """Posterior mean and sd of the latent yield surface (standardised units) on all 1,728 cells."""
    Xs = IDX[sel]
    L = np.linalg.cholesky(kern(h, d2, Xs, Xs) + (h["sigma"] ** 2 + 1e-6) * np.eye(len(sel)))
    Kgx = kern(h, d2, IDX, Xs)
    alpha = np.linalg.solve(L.T, np.linalg.solve(L, ys - h["c"]))
    V = np.linalg.solve(L, Kgx.T)
    sd = np.sqrt(np.maximum(prior_var(h) - np.sum(V**2, 0), 1e-12))
    return h["c"] + Kgx @ alpha, sd, (L, Kgx)


def prior_surface(h, d2, rng, k=1):
    """Draws of the surface on the whole grid. The interaction kernel is a Kronecker product of five
    small matrices, so a draw needs five tiny Cholesky factors instead of one 1,728 x 1,728 one."""
    R, M = level_mats(h, d2)
    LR = [np.linalg.cholesky(r + 1e-9 * np.eye(len(r))) for r in R]
    LM = [np.linalg.cholesky(r + 1e-9 * np.eye(len(r))) for r in M]
    out = np.empty((k, N))
    for j in range(k):
        z = rng.standard_normal(NL)
        for f in range(5):
            z = np.moveaxis(np.tensordot(LR[f], z, axes=([1], [f])), 0, f)
        s = h["eta"] * z.reshape(-1)
        if "amp" in h:
            for f in range(5):
                s += h["amp"][f] * (LM[f] @ rng.standard_normal(NL[f]))[IDX[:, f]]
        out[j] = h["c"] + s
    return out


def posterior_surface(h, d2, sel, ys, cache, rng, k=1):
    """Posterior draws by Matheron's rule: a prior draw, corrected by what it got wrong at the data."""
    L, Kgx = cache
    fp = prior_surface(h, d2, rng, k)
    eps = h["sigma"] * rng.standard_normal((k, len(sel)))
    resid = ys[None, :] - fp[:, sel] - eps
    a = np.linalg.solve(L.T, np.linalg.solve(L, resid.T))
    return fp + (Kgx @ a).T

T_BUDGET, N0 = 50, 5


def standardise(yv):
    s = yv.std()
    return (yv - yv.mean()) / (s if s > 0 else 1.0), yv.mean(), (s if s > 0 else 1.0)


class MapFitter:
    """Compile the PyMC model's log posterior and gradient ONCE; each refit swaps the data with
    pm.set_data (the arrays may change length) and runs L-BFGS on the unconstrained parameters."""

    def __init__(self, model, d2):
        self.m, self.d2 = model, d2
        self.fn = model.logp_dlogp_function(ravel_inputs=True)
        self.fn.set_extra_values({})
        ip = model.initial_point()
        rv = DictToArrayBijection.map({v.name: ip[v.name] for v in model.value_vars})
        self.x0, self.info = rv.data, rv.point_map_info
        self.x = self.x0.copy()

    def fit(self, sel, ys):
        pm.set_data({"idx": IDX[sel], "y": ys}, model=self.m)
        best = None
        for s in [self.x, self.x0]:                 # warm start, and the prior centre
            r = optimize.minimize(lambda q: tuple(-v for v in self.fn(q)), s, jac=True,
                                  method="L-BFGS-B")
            if best is None or r.fun < best.fun:
                best = r
        self.x = best.x
        return self.unpack(best.x)

    def unpack(self, x):
        h = {}
        for k, v in DictToArrayBijection.rmap(RaveledVars(x, self.info)).items():
            if k.endswith("_logodds__"):
                h[k[:-10]] = 1 / (1 + np.exp(-v))
            elif k.endswith("_log__"):
                h[k[:-6]] = np.exp(v)
            else:
                h[k] = v
        return h


def expected_improvement(mu, sd, best):
    z = (mu - best) / sd
    return (mu - best) * stats.norm.cdf(z) + sd * stats.norm.pdf(z)


def choose(acq, h, d2, sel, ys, mu, sd, cache, tested, rng, q=1):
    """The next q reactions. For q > 1: Thompson draws, or 'kriging believer' (pretend the
    posterior mean was observed, update, choose again)."""
    new = []
    if acq == "ts":
        draws = posterior_surface(h, d2, sel, ys, cache, rng, k=3 * q)
        for d in draws:
            d[tested] = -np.inf
            d[new] = -np.inf
            if np.argmax(d) not in new:
                new.append(int(np.argmax(d)))
            if len(new) == q:
                break
        return new
    s2, y2, mu2, sd2 = list(sel), ys.copy(), mu, sd
    for j in range(q):
        if acq == "ei":
            a = expected_improvement(mu2, sd2, ys.max())
        elif acq == "pi":
            a = stats.norm.cdf((mu2 - ys.max() - 0.01) / sd2)
        elif acq == "ucb":
            a = mu2 + 2.0 * sd2
        elif acq == "greedy":
            a = mu2.copy()
        a[tested] = -np.inf
        a[new] = -np.inf
        i = int(np.argmax(a))
        new.append(i)
        if j < q - 1:
            s2, y2 = s2 + [i], np.append(y2, mu2[i])
            mu2, sd2, _ = gp_posterior(h, d2, np.array(s2), y2)
    return new


def campaign(fitter, seed, acq="ei", q=1, T=T_BUDGET, noise_sd=0.0, incumbent="observed",
             keep_frames=False):
    """One optimisation campaign on the real table. Returns the chosen reactions in order, the
    yields seen, the largest EI in yield points before each choice, and the TRUE yield of the
    reaction the model would recommend (highest posterior mean among those run) at each step."""
    rng_c = np.random.default_rng(seed)
    sel = list(rng_c.choice(N, N0, replace=False))           # the same 5 starting reactions for
    y_obs = list(Y[sel] + noise_sd * rng_c.standard_normal(N0))  # every strategy with this seed
    fitter.x = fitter.x0.copy()
    out = {"max_ei": [], "rec": [], "frames": [], "hyper": []}
    while len(sel) < T:
        ys, m_y, s_y = standardise(np.array(y_obs))
        h = fitter.fit(np.array(sel), ys)
        mu, sd, cache = gp_posterior(h, fitter.d2, np.array(sel), ys)
        tested = np.zeros(N, bool)
        tested[sel] = True
        inc = ys.max() if incumbent == "observed" else mu[sel].max()
        e = expected_improvement(mu, sd, inc) * s_y
        e[tested] = 0.0
        out["max_ei"].append(e.max())
        out["rec"].append(Y[sel[int(np.argmax(mu[sel]))]])
        out["hyper"].append(h)
        if acq == "ei" and incumbent == "mean":        # EI against the best posterior mean (q = 1)
            new = [int(np.argmax(np.where(tested, -np.inf, e)))]
        else:
            new = choose(acq, h, fitter.d2, np.array(sel), ys, mu, sd, cache, tested, rng_c,
                         min(q, T - len(sel)))
        if keep_frames:
            out["frames"].append(dict(n=len(sel), mu=mu * s_y + m_y, sd=sd * s_y, ei=e, new=new,
                                      sel=list(sel)))
        sel += new
        y_obs += list(Y[new] + noise_sd * rng_c.standard_normal(len(new)))
    out.update(sel=np.array(sel[:T]), y_obs=np.array(y_obs[:T]))
    return out


fit_add = MapFitter(*gp_model("additive"))
tic = time.time()
demo = campaign(fit_add, seed=RANDOM_SEED, keep_frames=True)
print(f"one 50-reaction campaign (45 refits + choices): {time.time() - tic:.1f} s")
print("best yield after 5, 10, 15, 20, 30, 50 reactions:",
      [float(np.max(Y[demo['sel'][:n]])) for n in [5, 10, 15, 20, 30, 50]])

fig, axes = plt.subplots(1, 4, figsize=(12.5, 4.2), dpi=72, width_ratios=[1, 1, 1, 0.8], layout="none")
fig.subplots_adjust(left=0.08, right=0.98, top=0.84, bottom=0.12, wspace=0.16)
fr0 = demo["frames"][0]
ims = [draw_space(axes[0], fr0["mu"], vmin=0, vmax=100, labels=True),
       draw_space(axes[1], fr0["sd"], cmap="magma", vmin=0, vmax=40, labels=False),
       draw_space(axes[2], fr0["ei"], cmap="inferno", vmin=0, vmax=5, labels=False)]
for ax in axes[:3]:
    for lab in ax.get_yticklabels():
        lab.set_fontsize(6)
titles = ["posterior mean yield (0-100%)", "posterior sd (0-40 points)", "expected improvement (0-5 points)"]
dots = [ax.plot([], [], "o", mfc="none", mec="w", ms=3.5, mew=0.8)[0] for ax in axes[:3]]
nxt = [ax.plot([], [], "s", mfc="none", mec="cyan", ms=7, mew=1.6)[0] for ax in axes[:3]]
best_curve = np.maximum.accumulate(Y[demo["sel"]])
axes[3].plot(np.arange(1, T_BUDGET + 1), best_curve, color="0.7")
cur, = axes[3].plot([], [], color="C0", lw=2)
axes[3].axhline(Y.max(), color="k", ls=":", lw=0.8)
axes[3].set(xlim=(0, T_BUDGET), ylim=(0, 105), xlabel="reactions run", title="best yield so far")
axes[3].title.set_fontsize(9)


def update(k):
    fr = demo["frames"][k]
    for im, v in zip(ims, [fr["mu"], fr["sd"], fr["ei"]]):
        im.set_data(to_grid(v))
    r, c = grid_pos(np.array(fr["sel"]))
    rn, cn = grid_pos(np.array(fr["new"]))
    for d_, n_ in zip(dots, nxt):
        d_.set_data(c, r)
        n_.set_data(cn, rn)
    for ax, t in zip(axes, titles):
        ax.set_title(t, fontsize=8, pad=14)
    cur.set_data(np.arange(1, fr["n"] + 1), best_curve[: fr["n"]])
    fig.suptitle(f"after {fr['n']} reactions: best {best_curve[fr['n'] - 1]:.0f}%, "
                 f"largest EI {fr['ei'].max():.1f} points; cyan = reaction {fr['n'] + 1}", fontsize=10)
    return ims + dots + nxt + [cur]


FRAMES = list(range(0, 20)) + list(range(21, len(demo["frames"]), 2))   # every step to 25, then every other
anim = animation.FuncAnimation(fig, update, frames=FRAMES, interval=500)
plt.close(fig)
_fmt = plt.rcParams["animation.frame_format"]
plt.rcParams["animation.frame_format"] = "jpeg"     # heatmaps: JPEG frames are far smaller than PNG
display(HTML(anim.to_jshtml(default_mode="once")))
plt.rcParams["animation.frame_format"] = _fmt
seq = df.loc[demo["sel"][:24], FACT + ["yield"]].assign(n=np.arange(1, 25)).set_index("n")
print(seq.to_string())

NCAMP = 50
SEEDS = 1000 + np.arange(NCAMP)


def run_many(fitter, **kw):
    return [campaign(fitter, s, **kw) for s in SEEDS]


def best_curves(runs):
    return np.array([np.maximum.accumulate(Y[r["sel"]]) for r in runs])


def run_random(seed):
    rng_c = np.random.default_rng(seed)
    first = rng_c.choice(N, N0, replace=False)
    rest = np.setdiff1d(np.arange(N), first)
    return np.r_[first, rng_c.permutation(rest)[: T_BUDGET - N0]]


def run_ofat(seed):
    """One factor at a time from the best of the same 5 starting reactions: try every level of the
    ligand, keep the best, then base, solvent, temperature, concentration; repeat the cycle."""
    rng_c = np.random.default_rng(seed)
    sel = list(rng_c.choice(N, N0, replace=False))
    while len(sel) < T_BUDGET:
        grew = False
        for f in [0, 1, 2, 4, 3]:
            cur = IDX[sel[int(np.argmax(Y[sel]))]].copy()
            for lv in range(NL[f]):
                c = cur.copy()
                c[f] = lv
                i = int(np.ravel_multi_index(tuple(c), NL))
                if i not in sel and len(sel) < T_BUDGET:
                    sel.append(i)
                    grew = True
        if not grew:
            sel.append(int(rng_c.choice(np.setdiff1d(np.arange(N), sel))))
    return np.array(sel)


tic = time.time()
fit_prod = MapFitter(*gp_model("product"))
RUNS = {"BO: EI, additive kernel": run_many(fit_add, acq="ei"),
        "BO: EI, product kernel": run_many(fit_prod, acq="ei")}
print(f"2 x {NCAMP} BO campaigns: {time.time() - tic:.0f} s")
CURVES = {k: best_curves(v) for k, v in RUNS.items()}
CURVES["random search"] = np.array([np.maximum.accumulate(Y[run_random(s)]) for s in SEEDS])
CURVES["one factor at a time"] = np.array([np.maximum.accumulate(Y[run_ofat(s)]) for s in SEEDS])

game = data.load("edbo_arylation_game")
data.describe("edbo_arylation_game")
H = np.full((game.participant.nunique(), 100), np.nan)
for p, g in game.groupby("participant"):
    H[p, : len(g)] = np.maximum.accumulate(g.sort_values("step")["yield"].to_numpy())
n_play = np.sum(~np.isnan(H), 0)
H_carry = pd.DataFrame(H).ffill(axis=1).to_numpy()      # a chemist who stopped keeps their best
edbo = data.load("edbo_arylation_edbo_runs")
E = np.maximum.accumulate(edbo.pivot(index="run", columns="step", values="yield").to_numpy(), axis=1)
print(f"chemists: {len(H)}, reactions each: median {np.median(n_play.sum() / len(H)):.0f} "
      f"(still playing at 10/20/30/50: {n_play[9]}, {n_play[19]}, {n_play[29]}, {n_play[49]})")


def summarise(C, ns=(10, 20, 30, 50)):
    row = {}
    for n in ns:
        if n > C.shape[1]:
            continue
        v = C[:, n - 1]
        v = v[~np.isnan(v)]
        k = np.sum(v >= 95)
        lo, hi = stats.beta(1 + k, 1 + len(v) - k).ppf([0.05, 0.95])
        row[f"median best @{n}"] = np.median(v)
        row[f"P(>=95%) @{n}"] = f"{k / len(v):.2f} [{lo:.2f}, {hi:.2f}]"
    return row


tab = {"random search (exact)": {**{f"median best @{n}": med_random[n - 1] for n in (10, 20, 30, 50)},
                                 **{f"P(>=95%) @{n}": f"{p95_random[n - 1]:.2f}" for n in (10, 20, 30, 50)}}}
for k, C in CURVES.items():
    tab[k] = summarise(C)
tab["chemists (if they had stopped: best kept)"] = summarise(H_carry)
tab["EDBO's own simulations (batches of 5)"] = summarise(E)
TAB = pd.DataFrame(tab).T
print(TAB[[c for c in TAB.columns if c.startswith("median")]].round(1).to_string())
print(TAB[[c for c in TAB.columns if c.startswith("P(")]].to_string())