C13 · Scale-up from a partial screen — Bayesian Modelling in Practice

The brief

A process chemistry group has to fix the conditions for scaling up a C-H arylation. There are 1,728 candidate conditions (ligand, base, solvent, concentration, temperature). Their high-throughput robot ran a seeded random 60 of them last month. The head of the group asks:

"Our best screened condition gave about 91%. Do we lock it in for the pilot plant now, or do we buy another round of screening first? If another round, how many reactions? I want an expected value and the probability that it pays off - not a feeling."

The economics, agreed with finance:

every percentage point of yield at scale is worth $5,000 over the planned campaigns (assume screening yields carry over to the plant);
one more screening reaction costs $1,500** (materials, analysis, chemist time), and any extra round delays the pilot plant by a week, which costs **$10,000 whatever its size;
only a condition that has actually been run can go to the plant: after an extra round they take the best measured condition (screen or new round). Treat recorded yields as exact - the table has no replicates.

Your job: build a model of yield over the whole grid from the 60 reactions, say which conditions could be the best and how good the best could be, then turn that into the decision. Because this dataset is a complete screen, you can then do what no real group can: audit every prediction and every decision against the truth.

This is the sequel to example E82 (Bayesian optimisation). E82 is about the sequential loop (which experiment next, one at a time) and uses a Gaussian process on raw yields. Here the loop is not the point: the questions are what the model of a partial screen should look like and how to price the next batch. Do E82 first; its "Try it yourself" 1 and 2 are where this starts.

How this notebook works

Each task states what to deliver, not how. Write your code in the YOUR CODE HERE cells.
Stuck? h.hint("task3") reveals hints one level at a time: nudge → approach → code skeleton. Try to get by on nudges.
h.check("task2", sigma=...) compares your numbers with the reference solution.
A full worked solution lives in notebooks/solutions/. Open it only when you are done (or truly stuck).

import logging
import time

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import xarray as xr
from scipy.special import expit, logit

from pymc_challenges import Hints, data

RANDOM_SEED = 2021
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)  # several fits: no sampler banner per fit
T_START = time.time()

h = Hints("C13")
h.tasks()

C13 - Scale-up from a partial screen: model the yields, then price the next experiment

task0 - Look at the screen (2 hints)
task1 - A Gaussian ANOVA on raw yield (3 hints, 3 check(s))
task2 - A likelihood that knows what a yield is (3 hints, 2 check(s))
task3 - Interactions from 60 reactions (3 hints, 2 check(s))
task4 - The whole grid (3 hints, 3 check(s))
task5 - Commit now, or buy k more reactions? (3 hints, 3 check(s))
task6 - Replay the decision (3 hints, 3 check(s))

The data

The full table is loaded once and put in grid order (ligand x base x solvent x concentration x temperature, as in E82), so row i of IDX gives the level indices of condition i. The screen is the seeded random 60. TRUTH holds all 1,728 yields: do not use it before Task 4; from then on it is your audit, never an input to a model or a decision.

data.describe("edbo_arylation")
grid = data.load("edbo_arylation")
FACT = ["ligand", "base", "solvent", "concentration", "temperature"]
LEV = [sorted(grid[f].unique()) for f in FACT]
NL = [len(lv) for lv in LEV]
for f, lv in zip(FACT, LEV):
    grid[f + "_i"] = pd.Index(lv).get_indexer(grid[f])
grid = grid.sort_values([f + "_i" for f in FACT]).reset_index(drop=True)
IDX = grid[[f + "_i" for f in FACT]].to_numpy()
N = len(grid)
assert N == np.prod(NL) and (np.ravel_multi_index(IDX.T, NL) == np.arange(N)).all()

SCREEN = np.sort(np.random.default_rng(13).choice(N, 60, replace=False))  # the robot's 60 runs
screen = grid.loc[SCREEN, FACT + ["yield"]]
TRUTH = grid["yield"].to_numpy()  # the audit: hands off until Task 4
CANDIDATES = np.setdiff1d(np.arange(N), SCREEN)  # the 1,668 conditions nobody has run
for f, lv in zip(FACT, LEV):
    print(f"{f:13s} {len(lv):2d} levels: {', '.join(map(str, lv))}")
screen.sort_values("yield", ascending=False).head(8)

edbo_arylation: A fully enumerated high-throughput screen of a palladium-catalysed direct (C-H) arylation: every combination of 12 phosphine ligands x 4 carboxylate bases x 4 solvents x 3 concentrations x 3 temperatures, one yield each (1,728 rows): ligand, base, solvent, concentration (M: 0.057, 0.1, 0.153), temperature (C: 90, 105, 120), yield (%, 0-100).
  source: Shields, Stevens, Li, Parasram, Damani, Martinez Alvarado, Janey, Adams & Doyle (2021), Bayesian reaction optimization as a tool for chemical synthesis, Nature 590:89-96; data from the authors' repository github.com/b-shields/edbo (MIT licence, (c) 2020 Benjamin J. Shields). BUILT by tools/build_e82_arylation.py (SMILES replaced by the names in the repository's *-list.csv files) - keep the cached file.
  url:    https://raw.githubusercontent.com/b-shields/edbo/master/experiments/data/direct_arylation/experiment_index.csv
ligand        12 levels: BrettPhos, CgMe-PPh, GorlosPhos HBF4, JackiePhos, P(fur)3, PCy3 HBF4, PPh2Me, PPh3, PPhMe2, PPhtBu2, X-Phos, tBPh-CPhos
base           4 levels: CsOAc, CsOPiv, KOAc, KOPiv
solvent        4 levels: BuCN, BuOAc, DMAc, p-Xylene
concentration  3 levels: 0.057, 0.1, 0.153
temperature    3 levels: 90, 105, 120

	ligand	base	solvent	concentration	temperature	yield
202	CgMe-PPh	CsOPiv	DMAc	0.100	105	91.11
199	CgMe-PPh	CsOPiv	DMAc	0.057	105	91.06
1452	X-Phos	CsOAc	BuOAc	0.100	90	85.24
1444	X-Phos	CsOAc	BuCN	0.100	105	83.83
1678	tBPh-CPhos	KOAc	DMAc	0.100	105	78.95
1534	X-Phos	KOAc	DMAc	0.100	105	65.37
1525	X-Phos	KOAc	BuOAc	0.100	105	64.62
4	BrettPhos	CsOAc	BuCN	0.100	105	64.48

Task 0 · Look at the screen

Only the 60 screened yields. Deliver

The distribution of the yields: how many are exactly 0, how many are traces (below 1%), how close does anything come to 100%?
How much of the design does the screen cover: reactions per ligand, and how many of the 48 ligand x solvent and 48 ligand x base combinations were run at all?
One sentence: what does a likelihood for these numbers have to be able to do?

Your turn Work this task in the notebook.

Task 1 · The default: a Gaussian ANOVA on raw yield

The model most people write first: yield (%) = intercept + one effect per level of each of the five factors + Normal noise, with the effects of each factor shrunk towards zero by a shared scale (a hierarchical ANOVA, E02).

Deliver

A fitted model with clean diagnostics.
A posterior predictive check aimed at the features you listed in Task 0.
Predictions for all 1,728 conditions (the "table" the model expects - latent mean plus noise): what share of them is negative, and what is the posterior of the best achievable yield (the largest predicted yield over the grid)? Would you show this to the head of the group?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Write the model as intercept + one zero-sum effect vector per factor, each with its own scale. With only 3 levels for concentration and temperature, how well can 60 points pin down those scales - and what does that mean for centred vs non-centred effects? For the PPC, pick statistics that the Gaussian is not free to fit: the pile at zero, the maximum.

2 Approach

pm.ZeroSumNormal("z_f", sigma=1.0, dims=f) times a pm.HalfNormal scale, indexed by the level codes of the 60 screened rows. Raw-scale priors: effects of tens of points are plausible (HalfNormal(20)). For the grid, push every posterior draw through all 1,728 rows of IDX in NumPy (post["e_ligand"].transpose("ligand", "sample").values[IDX[:, 0]]), add sigma * standard_normal, and take the share below 0 and the maximum per draw.

3 Skeleton

with pm.Model(coords=COORDS) as naive:
    mu = pm.Normal("a", 30.0, 20.0)
    for j, f in enumerate(FACT):
        s = pm.HalfNormal(f"s_{f}", ___)
        e = pm.Deterministic(f"e_{f}", s * pm.ZeroSumNormal(f"z_{f}", sigma=1.0, dims=f), dims=f)
        mu = mu + e[X[:, j]]
    pm.Normal("y", mu, pm.HalfNormal("sigma", 20.0), observed=y)
    idata = pm.sample(target_accept=0.95, random_seed=RANDOM_SEED)

post = az.extract(idata, num_samples=2000)
mu = post["a"].values[None, :] + sum(
    post[f"e_{f}"].transpose(f, "sample").values[IDX[:, j]] for j, f in enumerate(FACT))
table = mu.T + post["sigma"].values[:, None] * rng.standard_normal(___)   # (draws, 1728)
share_negative = np.mean(table < 0)
best_achievable = np.median(table.max(axis=___))

Self-checks in the notebook (h.check("task1", …)):

sigma — posterior mean of the residual sd on the raw scale, about 17.5 yield points
share_negative — share of all predicted yields (draws x 1,728 conditions) below 0, about 0.22
best_achievable — posterior median of the largest predicted yield over the grid, about 113% (sic)

Task 2 · A likelihood that knows what a yield is

Deliver

A likelihood that respects the bounds of a yield and produces a pile of zeros/traces and a ceiling naturally. Justify the scale on which factors act, and how you treat 0 and 100.
A prior predictive check on the yield scale: what share of reactions does your prior expect to fail, and what share to exceed 95%?
The fit, diagnostics, and the same posterior predictive check as in Task 1.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Two separate problems: the outcome is bounded, and a third of the reactions fail outright. On which scale would "this ligand doubles the odds of conversion" be additive? And is the difference between 0% and 0.3% something you want a model to explain?

2 Approach

Model logit(yield / 100) as Normal, and treat the ends as censored: anything below 1% means "trace or nothing" (latent below logit(0.01)), anything above 99.5% means "quantitative". pm.Censored(pm.Normal.dist(mu, sigma), lower=..., upper=..., observed=logit(np.clip(y, 1, 99.5) / 100)) does exactly that. A Tobit model on the raw scale (censored at 0 and 100) also respects the bounds, but its noise has the same spread at 85% as at 40% - remember that when you get to the decision.

Priors now live on the logit scale: an intercept Normal(-1, 1.5) puts a typical reaction around 25%; effect scales HalfNormal(1.5). Check them with pm.sample_prior_predictive transformed back with expit.

3 Skeleton

Z_LO, Z_HI = logit(0.01), logit(0.995)
with pm.Model(coords=COORDS) as additive:
    mu = pm.Normal("a", -1.0, 1.5)
    for j, f in enumerate(FACT):
        ...                                                   # as in Task 1, HalfNormal(1.5)
    sigma = pm.HalfNormal("sigma", ___)
    pm.Censored("y", pm.Normal.dist(mu, sigma), lower=___, upper=___,
                observed=logit(np.clip(y, 1.0, 99.5) / 100))
    prior = pm.sample_prior_predictive(500)
    idata = pm.sample(target_accept=0.95)
    pm.sample_posterior_predictive(idata, extend_inferencedata=True)
rep = 100 * expit(az.extract(idata, group="posterior_predictive", var_names=["y"]).values.T)
share_trace = np.mean(rep < 1.0 + 1e-9)                       # censored draws sit at the bound

Self-checks in the notebook (h.check("task2", …)):

sigma — posterior mean of the residual sd on the LOGIT scale for the additive censored model, about 1.7
share_trace — share of posterior-predictive screen yields below 1%, about 0.27 (observed 0.30)

Task 3 · Interactions from 60 reactions

Chemists will tell you that a ligand can work in one solvent and fail in another. With 60 reactions spread over 48 ligand x solvent cells, most cells have been run once or never.

Deliver

Add ligand x solvent, ligand x base and base x solvent interactions to the model of Task 2, each shrunk towards zero by its own scale. Read the diagnostics carefully: something will be harder to sample than before. Find which parameters, explain why in terms of the data, and decide whether it matters for what the group needs - predictions of yields.
Posterior of the three interaction scales, and a LOO comparison with the additive model.
Which ligand x solvent combinations does the model think are special, and how sure is it?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

48 ligand x solvent effects from 60 reactions: almost all of what you learn about them comes through their shared scale. When the data say little about a group of effects and their scale together, what geometry does the posterior have - and which parameterisation is built for that?

2 Approach

pm.ZeroSumNormal(..., dims=("ligand", "solvent"), n_zerosum_axes=2) keeps rows and columns summing to zero, so interactions do not duplicate main effects (E36). The centred version ZeroSumNormal(sigma=s), and the non-centred one (a unit ZeroSumNormal times s) both struggle: plot sigma against each interaction scale. If two parameters trade off because the data cannot tell them apart, no reparameterisation removes that - check r_hat and ESS of what you will use (the latent mean at the promising conditions). Compare models with pm.compute_log_likelihood + az.compare - both have the same censored likelihood, so LOO is a fair comparison here (it would not be between the raw and logit models).

3 Skeleton

for p, q in [(0, 2), (0, 1), (1, 2)]:
    name, dims = f"{FACT[p]}_x_{FACT[q]}", (FACT[p], FACT[q])
    s = pm.HalfNormal(f"s_{name}", 1.0)
    e = pm.Deterministic(f"e_{name}", s * pm.ZeroSumNormal(
        f"z_{name}", sigma=1.0, dims=dims, n_zerosum_axes=___), dims=dims)
    mu = mu + e[X[:, p], X[:, q]]
...
cmp = az.compare({"additive": idata_add, "interactions": idata_int}, round_to=1)
elpd_diff = cmp.loc["interactions", "elpd"] - cmp.loc["additive", "elpd"]

Self-checks in the notebook (h.check("task3", …)):

s_ligand_solvent — posterior mean of the ligand x solvent interaction scale (logit), about 1.05
elpd_diff — elpd(interactions) - elpd(additive), about +10 nats (with many Pareto k > 0.7)

Task 4 · The whole grid: which condition is best, and how good can it get?

Deliver, from your best model (and, for contrast, the Task 1 model):

For every one of the 1,728 conditions, the posterior probability that it is the best one (screened conditions keep their measured yield). The ten most likely, with their posterior yield. How concentrated is the probability?
The posterior of the best achievable yield and of the probability that some condition reaches 95%.
Now open the audit: score each model's predictions of the 1,668 unscreened yields against TRUTH (a proper scoring rule, interval coverage overall and among the truly good conditions), and say where the probability of being best actually went.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

"Probability that condition c is the best" is a statement about the joint posterior of all 1,728 yields: in each posterior draw exactly one condition wins. What do you do with the 60 you have already measured, and with draws that tie at 100%?

2 Approach

From each posterior draw build one plausible table (latent mean + noise), overwrite the screened conditions with their measured values, and take argmax per draw (on the logit scale, where nothing ties); np.bincount(argmax, minlength=1728) / n_draws is P(best). The maximum per draw on the yield scale is the best achievable yield. For the audit, the sample CRPS is mean|Y - y| - 0.5 mean|Y - Y'| (sort the draws to compute the second term in O(S log S)); compare 90% interval coverage overall and among conditions that are truly at least 80%.

3 Skeleton

lat = G["latent"].copy()                       # (draws, 1728), logit scale
lat[:, SCREEN] = logit(np.clip(y, 1, 99.5) / 100)
p_best = np.bincount(lat.argmax(axis=1), minlength=N) / len(lat)
yy = G["yield"].copy(); yy[:, SCREEN] = y
best_achievable = yy.max(axis=___)

def crps(draws, obs):
    S = draws.shape[0]
    s = np.sort(draws, 0)
    w = (2 * np.arange(1, S + 1) - S - 1)[:, None]
    return np.abs(draws - obs).mean(0) - (w * s).sum(0) / S**2
crps(G["yield"][:, CANDIDATES], TRUTH[CANDIDATES]).mean()

Self-checks in the notebook (h.check("task4", …)):

p_best_top — the largest P(best) of any condition under the interaction model, about 0.02
p_reach_95 — posterior probability that some condition reaches 95% (interaction model), nearly 1
crps — mean CRPS (yield points) of the interaction model on the 1,668 held-out yields, about 7.9

Task 5 · Commit now, or buy k more reactions?

Deliver

For k = 0, 4, 8, 12, 16, 24, 32, 48 extra reactions: which k conditions you would run (your rule, from the posterior), the posterior distribution of the gain in yield over the current best measured condition, the expected net value in dollars (brief's economics), and the probability that it pays (net value > 0).
The same table from the Task 1 model. Do the two models recommend the same thing?
A recommendation to the head of the group in three sentences.
Then, using TRUTH: what would each recommendation actually have earned?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

After the round, the plant gets the best measured condition. So the value of a batch is the expected value of the maximum of its yields - not the sum of each reaction's expected improvement. Which conditions make a good batch, and why might the posterior's noise term matter more here than anywhere else?

2 Approach

Use the grid draws: for a batch B, gain per draw = max(0, max_{j in B} y_j - 91.1). Build B greedily: start from the incumbent, add the unscreened condition that most raises the mean of max(current, y_c) over draws, repeat to 48 (a Monte-Carlo q-EI). Nested batches let you read off every k from one ordering. Then net = 5000 * gain - 1500 * k - 10000 * (k > 0), per draw: mean, interval, P(net > 0). Do the same with the Task 1 draws.

3 Skeleton

cur = np.full(n_draws, 91.11); pool = list(CANDIDATES); batch = []
for _ in range(48):
    score = np.maximum(yld[:, pool], cur[:, None]).mean(axis=0)
    c = pool.pop(int(np.argmax(score)))
    batch.append(c); cur = np.maximum(cur, yld[:, c])
for k in [0, 4, 8, 12, 16, 24, 32, 48]:
    g = np.maximum(yld[:, batch[:k]].max(axis=1) - 91.11, 0) if k else np.zeros(n_draws)
    net = 5000 * g - 1500 * k - 10000 * (k > 0)
    print(k, g.mean(), net.mean(), np.mean(net > ___))

Self-checks in the notebook (h.check("task5", …)):

expected_gain_8 — expected gain (yield points) of the best 8-reaction batch under the interaction model, about 5.8
best_k — the k with the highest expected net value under the interaction model - a small round (4 or 8)
p_pays — P(net value > 0) at the recommended k (at k = 8 if the recommendation is 0), about 0.7

Task 6 · One screen is one anecdote: replay the decision

A single realised outcome cannot tell you whether a decision rule is good. Because you have the full table, you can rerun the group's situation on other random screens of 60.

Deliver for at least six other screens, with both the Task 1 model and your best model: the recommended k, the expected net value, and the realised net value against TRUTH. Then a calibration check: where does the realised gain fall in each model's predictive distribution of the gain (a PIT value)? Is either model systematically over- or under-optimistic?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Wrap everything from the screen to the recommended k in one function of the screen's seed. Then ask two different questions: did the recommendations make money on average, and was each model's predicted gain distribution honest about the realised one?

2 Approach

For each seed: draw 60 conditions, fit both models, build the batch, take the recommended k, and record the expected net value and the realised net value from TRUTH. A PIT value is mean(g_draws < g_true) + 0.5 * mean(g_draws == g_true) (the mid-rank handles the point mass at a gain of 0). PITs near 0 mean the model promised more than it delivered. Six screens x two models is about two minutes; del each idata.

3 Skeleton

def pipeline(seed):
    scr = np.sort(np.random.default_rng(seed).choice(N, 60, replace=False))
    ys, Xs, inc = TRUTH[scr], IDX[scr], TRUTH[scr].max()
    cand = np.setdiff1d(np.arange(N), scr)
    for scale, pairs in [("raw", ()), ("logit", PAIRS)]:
        yld = ___                                     # fit on (Xs, ys), grid draws
        batch = greedy_batch(yld, 48, inc, cand)
        ...                                           # value table, k*, realised, PIT

Self-checks in the notebook (h.check("task6", …)):

mean_pit_logit — mean PIT of the realised gain under the interaction model over screens 1-6, about 0.8 (conservative)
realised_naive — mean realised net value ($ thousands) of the Gaussian model's recommendations, about 16
realised_logit — mean realised net value ($ thousands) of the interaction model's recommendations, about 45

Going further

Correlated residuals. The interaction model's noise is independent across conditions, so a batch of near-identical conditions looks like several independent lottery tickets. Add a ligand x base x solvent effect (or a GP over concentration and temperature within each cell, E82) and see whether the batch spreads out and the value of a round changes.
Batches that learn. The greedy batch ignores that running a reaction also teaches the model about its neighbours. Compare it with two rounds of half the size, refitting in between (a two-stage decision; E82's loop is the limit).
The chemists' game. edbo_arylation_game holds 50 chemists' choices on the same table. Give each chemist's first 60 reactions to the model as the screen: is a chemist's (non-random) screen worth more or less to the model than a random one, and does the decision change?
Noisy yields. Add replicate noise (say sd 5 points) to the screen and to the extra round, and let the group re-run the apparent winner before committing (E82 part H). What is a replicate worth?

h.progress()

C13: 0/20 hints used, 16/16 checks passed.

Reference solution

A complete, executed solution with commentary on every task — including the ways the obvious model fails. Try the hint ladder first.

Open the worked solution →

#The brief

#How this notebook works

#The data

#Task 0 · Look at the screen

#Task 1 · The default: a Gaussian ANOVA on raw yield

#Task 2 · A likelihood that knows what a yield is

#Task 3 · Interactions from 60 reactions

#Task 4 · The whole grid: which condition is best, and how good can it get?

#Task 5 · Commit now, or buy k more reactions?

#Task 6 · One screen is one anecdote: replay the decision

#Going further

Reference solution

The brief

How this notebook works

The data

Task 0 · Look at the screen

Task 1 · The default: a Gaussian ANOVA on raw yield

Task 2 · A likelihood that knows what a yield is

Task 3 · Interactions from 60 reactions

Task 4 · The whole grid: which condition is best, and how good can it get?

Task 5 · Commit now, or buy k more reactions?

Task 6 · One screen is one anecdote: replay the decision

Going further