| """Offline unit checks for the Lecture 3 examples.""" |
| import csv |
| import itertools |
| from pathlib import Path |
| import sys |
| from types import SimpleNamespace |
| import unittest |
|
|
| import torch |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT / "lecture_3")) |
| import esm2_flow_guidance as flow |
| import esm2_diffusion_guidance as diffusion |
|
|
|
|
| class GuidanceTests(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| torch.set_num_threads(2) |
|
|
| def test_example_annotations(self): |
| with (ROOT / "lecture_3" / "esm2_example.csv").open(newline="") as handle: |
| rows = list(csv.DictReader(handle)) |
| self.assertEqual(len(rows), 64) |
| self.assertEqual({len(row["sequence"]) for row in rows}, {24}) |
| labels = torch.tensor([[float(row["r1"]), float(row["r2"])] for row in rows]) |
| for module in (flow, diffusion): |
| computed = module.composition_proxies([row["sequence"] for row in rows]) |
| torch.testing.assert_close(computed, labels, atol=1e-7, rtol=0) |
| for row in rows: |
| self.assertEqual(int(row["c"]), int(float(row["r1"]) > 0)) |
|
|
| def test_scalarization_weights(self): |
| for module in (flow, diffusion): |
| with self.subTest(module=module.__name__): |
| weights = module.normalize_weights((7, 3)).cpu() |
| torch.testing.assert_close(weights, torch.tensor([0.7, 0.3])) |
| for invalid in ((0, 0), (-1, 2), (1,), (float("nan"), 1)): |
| with self.assertRaises(ValueError): |
| module.normalize_weights(invalid) |
|
|
| def test_reward_gradient_without_weight_updates(self): |
| for module in (flow, diffusion): |
| with self.subTest(module=module.__name__): |
| torch.manual_seed(3) |
| reward = module.RewardModel(4, 5).to(module.DEVICE) |
| reward.eval().requires_grad_(False) |
| state = torch.randn(2, 4, 5, device=module.DEVICE) |
| time = torch.tensor([0.3, 0.7], device=module.DEVICE) |
| weights = module.normalize_weights((0.7, 0.3)) |
| gradient = module.reward_gradient(reward, state, time, weights) |
| delta = torch.zeros_like(state) |
| delta[0, 1, 2] = 1e-3 |
| difference = reward(state + delta, time) - reward(state - delta, time) |
| estimate = (difference * weights).sum() / (2e-3) |
| torch.testing.assert_close(gradient[0, 1, 2], estimate, atol=3e-5, rtol=3e-3) |
| self.assertFalse(gradient.requires_grad) |
| self.assertTrue(torch.isfinite(gradient).all()) |
| self.assertTrue(all(p.grad is None for p in reward.parameters())) |
|
|
| def test_ddpm_schedule(self): |
| self.assertEqual(diffusion.alpha_bars[0].item(), 1.0) |
| self.assertLess(diffusion.alpha_bars[-1].item(), 1e-4) |
| self.assertEqual(diffusion.posterior_variances[1].item(), 0.0) |
| self.assertTrue((diffusion.posterior_variances >= 0).all()) |
| self.assertTrue(torch.isfinite(diffusion.posterior_variances).all()) |
|
|
| def test_constrained_decoder_against_exhaustive_search(self): |
| torch.manual_seed(11) |
| scores = torch.randn(3, 4, 20) |
| tokenizer = SimpleNamespace(convert_tokens_to_ids=lambda tokens: list(range(20))) |
| esm = SimpleNamespace(lm_head=torch.nn.Identity()) |
| stats = {"z_mean": torch.zeros(1, 1, 20), "z_std": torch.ones(1, 1, 20)} |
| for module in (flow, diffusion): |
| polar = [i for i, a in enumerate(module.AMINO_ACIDS) if a in module.POLAR_RESIDUES] |
| nonpolar = [i for i in range(20) if i not in polar] |
| for minimum in range(5): |
| sequences = module.decode(scores, esm, tokenizer, stats, min_polar=minimum) |
| for batch, sequence in enumerate(sequences): |
| count = sum(a in module.POLAR_RESIDUES for a in sequence) |
| self.assertGreaterEqual(count, minimum) |
| actual = sum(float(scores[batch, j, module.AMINO_ACIDS.index(a)]) |
| for j, a in enumerate(sequence)) |
| optimal = max( |
| sum(float(scores[batch, j, polar if flag else nonpolar].max()) |
| for j, flag in enumerate(assignment)) |
| for assignment in itertools.product((0, 1), repeat=4) |
| if sum(assignment) >= minimum |
| ) |
| self.assertAlmostEqual(actual, optimal, places=5) |
| with self.assertRaises(ValueError): |
| module.decode(scores, esm, tokenizer, stats, min_polar=5) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
