<jupyter_start><jupyter_text>Predicting transcription factor binding affinity<jupyter_code>from collections import OrderedDict
from pprint import pprint
import pyro
import torch
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

from pyroed.datasets.data import load_tf_data
from pyroed.constraints import AllDifferent, Iff, IfThen, TakesValue
from pyroed.oed import thompson_sample
from pyroed.testing import generate_fake_data

matplotlib.rcParams["figure.facecolor"] = "white"<jupyter_output><empty_output><jupyter_text>Transcription factor dataWe consider data from [Survey of variation in human transcription factors reveals prevalent DNA binding changes](https://www.science.org/doi/abs/10.1126/science.aad2257),in particular for the PBX4 transcription factor.The dataset consists of measurements of the binding affinities of PBX4 to allpossible DNA sequences of length 8, i.e. for a total of $4^8 = 65536$ sequences.Since this dataset is exhaustive we can use it to do a semi-synthetic experiment in whichwe first "measure" a small number of binding affinities and then do additional "experiments" in multiple rounds.<jupyter_code>data = load_tf_data(data_dir="../pyroed/datasets")
for k, v in data.items():
    print(f"{k} {v.dtype} {tuple(v.shape)}")<jupyter_output>sequences torch.int64 (65792, 8)
responses torch.float32 (65792,)
batch_ids torch.int64 (65792,)<jupyter_text>Note that there are actually more than $65536$ data points due to some repeats. Sequences take values 0,1,2,3 corresponding to nucleotides.<jupyter_code>set(data["sequences"].reshape(-1).tolist())<jupyter_output><empty_output><jupyter_text>The response variable appears to be approximately Gaussian distributed.<jupyter_code>plt.hist(data["responses"].numpy(), bins=20)
plt.show()<jupyter_output><empty_output><jupyter_text>Modeling in pyroedSpecify the design space via SCHEMA, CONSTRAINTS, FEATURE_BLOCKS, and GIBBS_BLOCKS.<jupyter_code>SCHEMA = OrderedDict()
for n in range(8):
    SCHEMA[f"Nucleotide{n}"] = ["A", "C", "G", "T"]

CONSTRAINTS = []  # No constraints.

singletons = [[name] for name in SCHEMA]
pairs = [list(ns) for ns in zip(SCHEMA, list(SCHEMA)[1:])]
triples = [list(ns) for ns in zip(SCHEMA, list(SCHEMA)[1:], list(SCHEMA)[2:])]

FEATURE_BLOCKS = singletons + pairs
GIBBS_BLOCKS = triples<jupyter_output><empty_output><jupyter_text>Let's start with a random subsample of data.<jupyter_code>pyro.set_rng_seed(0)
full_size = len(data["responses"])
batch_size = 10
ids = torch.randperm(full_size)[:batch_size]
experiment = {k: v[ids] for k, v in data.items()}
print(f"Best response = {experiment['responses'].max():0.6g}")<jupyter_output>Best response = 0.574323<jupyter_text>Each step of the OED process we'll test on new data. Let's make a helper to simulate lab work.<jupyter_code>seq_to_id = {tuple(row): i for i, row in enumerate(data["sequences"].tolist())}

def update_experiment(experiment: dict, design: set) -> dict:
    batch_ids = experiment["batch_ids"].max().item() + 1
    ids = list(map(seq_to_id.__getitem__, sorted(design)))
    new_data = {
        "sequences": data["sequences"][ids],
        "batch_ids": torch.full((len(ids),), batch_ids),
        "responses": data["responses"][ids],
    }
    experiment = {k: torch.cat([v, new_data[k]]) for k, v in experiment.items()}
    print(f"Best response = {experiment['responses'].max():0.6g}")
    return experiment

def make_design(experiment: dict) -> set:
    return thompson_sample(
        SCHEMA,
        CONSTRAINTS,
        FEATURE_BLOCKS,
        GIBBS_BLOCKS,
        experiment,
        inference="svi",
        svi_num_steps=201,
        sa_num_steps=201,
        log_every=0,
        jit_compile=False,
    )<jupyter_output><empty_output><jupyter_text>Initialize our sequence of experiments:<jupyter_code>experiments = [experiment]<jupyter_output><empty_output><jupyter_text>Let's start with a single loop of active learning:<jupyter_code>%%time
design = make_design(experiments[-1])
print("Design:")
pprint(design)
experiments.append(update_experiment(experiments[-1], design))<jupyter_output>Design:
{(0, 2, 3, 3, 1, 2, 1, 1),
 (0, 3, 0, 2, 0, 1, 3, 1),
 (1, 0, 3, 1, 0, 1, 0, 0),
 (1, 0, 3, 3, 3, 3, 1, 0),
 (1, 3, 0, 2, 0, 0, 3, 0),
 (2, 1, 1, 0, 2, 3, 2, 1),
 (2, 1, 3, 1, 3, 3, 3, 1),
 (2, 2, 0, 0, 2, 3, 1, 1),
 (3, 1, 3, 0, 0, 0, 3, 0),
 (3, 3, 3, 0, 0, 2, 3, 1)}
Best response = 0.905897
CPU times: user 4.91 s, sys: 17.3 ms, total: 4.93 s
Wall time: 4.93 s<jupyter_text>Let's run multiple loops, say 10 more loops.<jupyter_code>%%time
for step in range(10):
    design = make_design(experiments[-1])
    experiments.append(update_experiment(experiments[-1], design))<jupyter_output>Best response = 0.905897
Best response = 0.905897
Best response = 0.905897
Best response = 0.94027
Best response = 0.94027
Best response = 0.959295
Best response = 0.959295
Best response = 0.959295
Best response = 0.959295
Best response = 0.959295
CPU times: user 53.3 s, sys: 175 ms, total: 53.4 s
Wall time: 53.5 s<jupyter_text>How did the response improve over time?<jupyter_code>responses = [e["responses"].max().item() for e in experiments]
plt.plot(responses)
plt.xlabel("Number of OED steps")
plt.ylabel("best binding affinity");<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13250950570342204, 'fraction_numerical': 0.05076045627376426, 'mean_word_length': 5.30059880239521, 'pattern_counts': {'":': 3, '<': 41, '<?xml version=': 0, '>': 43, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '205892', 'license': 'Apache-2.0', 'n_tokens_mistral': 2081, 'n_tokens_neox': 1858, 'n_words': 521, 'path': 'examples/tf8_demo.ipynb', 'repo_name': 'fritzo/pyroed', 'size': '5260'}