File size: 5,088 Bytes
c8293a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """CPU teaching utilities shared by the two lecture folders."""
import csv
import json
import random
from pathlib import Path
import numpy as np
import torch
from torch import nn
from lecture_core import DNA, K, MASK, encode
def seed_all(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.set_num_threads(1)
def decode(tokens):
alphabet = 'ACGTm'
return [''.join(alphabet[int(i)] for i in row) for row in tokens]
def make_data(count=512, length=8, seed=7):
"""Four synthetic motif families, with independent 8% base mutations."""
rng = np.random.default_rng(seed)
motifs = ['ACGT', 'CGTA', 'TATA', 'GCGC']
strings, labels = [], []
for _ in range(count):
motif = motifs[int(rng.integers(4))]
seq = list((motif * ((length + 3) // 4))[:length])
for j in range(length):
if rng.random() < .08:
seq[j] = 'ACGT'[int(rng.integers(4))]
strings.append(''.join(seq))
labels.append(int(sum(x in 'GC' for x in seq) / length >= .6))
return encode(strings), torch.tensor(labels)
def load_data(path, length):
if path is None:
return make_data(length=length)
with open(path) as f:
rows = list(csv.DictReader(f, delimiter='\t'))
strings = [r['sequence'].strip().upper() for r in rows]
if not strings or any(len(s) != length or set(s) - set('ACGT') for s in strings):
raise ValueError('All DNA sequences must contain only A/C/G/T and have --length bases.')
labels = [int(r.get('label', sum(c in 'GC' for c in s) / length >= .6))
for r, s in zip(rows, strings)]
if set(labels) - {0, 1}:
raise ValueError('Labels must be 0 or 1.')
return encode(strings), torch.tensor(labels)
class ConditionalDNA(DNA):
"""The slide network plus a label embedding; label 2 means unconditional."""
def __init__(self, width=32, max_len=64):
super().__init__(width, max_len)
self.condition = nn.Embedding(3, width)
def forward(self, z, t=None, label=None):
h = self.token(z) if z.ndim == 2 else self.soft(z)
pos = torch.arange(z.shape[1], device=z.device)
h = h + self.position(pos)[None]
if t is not None:
h = h + self.time(t[:, None])[:, None]
if label is None:
label = torch.full((len(z),), 2, dtype=torch.long, device=z.device)
h = h + self.condition(label)[:, None]
return self.output(self.context(h))
def objectives(tokens):
"""Both toy objectives are maximized: GC fraction and ATAT agreement."""
onehot = torch.nn.functional.one_hot(tokens.long(), 4).float()
return soft_objectives(onehot)
def soft_objectives(z):
gc = (z[..., 1] + z[..., 2]).mean(-1)
motif = torch.tensor([0, 3], device=z.device).repeat((z.shape[1] + 1) // 2)[:z.shape[1]]
match = z.gather(-1, motif[None, :, None].expand(z.shape[0], -1, 1)).squeeze(-1).mean(-1)
return torch.stack((gc, match), -1)
def metrics(tokens):
strings = decode(tokens)
score = objectives(tokens).mean(0)
return dict(valid_dna=all(set(s) <= set('ACGT') for s in strings),
unique_fraction=len(set(strings)) / len(strings),
mean_gc=float(score[0]), mean_atat_match=float(score[1]))
def save_run(out, config, losses, samples, extra=None):
out = Path(out)
out.mkdir(parents=True, exist_ok=True)
(out / ('sample_config.json' if config.get('mode') == 'sample' else 'config.json')).write_text(json.dumps(config, indent=2))
if losses or not (out / 'losses.json').exists():
(out / 'losses.json').write_text(json.dumps(losses, indent=2))
(out / 'samples.txt').write_text('\n'.join(decode(samples)) + '\n')
report = {'metrics': metrics(samples), **(extra or {})}
(out / 'report.json').write_text(json.dumps(report, indent=2))
print(json.dumps({'output': str(out), **report}, indent=2))
return report
def optimize(model, loss_fn, data, steps, batch_size=32, lr=.002):
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
losses = []
model.train()
for step in range(steps):
idx = torch.randint(len(data), (batch_size,))
optimizer.zero_grad(set_to_none=True)
loss = loss_fn(model, data[idx], idx)
if not torch.isfinite(loss):
raise FloatingPointError(f'Nonfinite loss at step {step}')
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
optimizer.step()
losses.append(float(loss.detach()))
model.eval()
return losses
def project_simplex(z, eps=1e-6):
"""Euclidean simplex projection, followed by a small interior floor."""
sorted_z = z.sort(-1, descending=True).values
cssv = sorted_z.cumsum(-1) - 1.
k = torch.arange(1, z.shape[-1] + 1, dtype=z.dtype, device=z.device)
active = sorted_z - cssv / k > 0
rho = active.sum(-1, keepdim=True).clamp_min(1)
theta = cssv.gather(-1, rho - 1) / rho
result = (z - theta).clamp_min(eps)
return result / result.sum(-1, keepdim=True)
|