File size: 3,343 Bytes
81ae663 | 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 | import json
from pathlib import Path
import numpy as np
import pytest
from dooable.chemistry import (
read_catalog,
reactions,
build_graph,
replay,
descriptor_rewards,
)
from dooable.graph import Graph, toy_graph
from dooable.exact import solve, sample, endpoint_distribution
from dooable.learning import train, load_model
from dooable.properties import scaffold_split
ROOT = Path(__file__).resolve().parents[1]
def test_chemical_route_replay_and_serialization(tmp_path):
ex = ROOT / "data/examples"
g = build_graph(
read_catalog(ex / "parents.csv"),
read_catalog(ex / "reagents.csv"),
reactions(ex / "reactions.json"),
2,
)
g.save(tmp_path / "graph.json")
gg = Graph.load(tmp_path / "graph.json")
assert len(gg.nodes) == len(g.nodes)
sol = solve(gg, descriptor_rewards(gg), 0.7)
rows = sample(gg, sol.forward, 200, 7)
assert all(replay(r, 2) for r in rows)
for y, v in gg.terminals.items():
assert gg.nodes[v].outcome == y
def test_neural_training_and_checkpoint(tmp_path):
g = toy_graph(4)
m, h = train(
g, {"A": 0.0, "B": 0.0}, steps=600, batch_size=32, seed=11, output=tmp_path
)
loaded, _ = load_model(tmp_path)
np.testing.assert_allclose(m.probabilities(), loaded.probabilities(), atol=1e-8)
assert h[-1]["endpoint_tv"] < 0.06
def test_scaffold_split_separation():
from rdkit.Chem.Scaffolds import MurckoScaffold
smi = [
"CC",
"CCC",
"CCO",
"c1ccccc1",
"Cc1ccccc1",
"c1ccncc1",
"C1CCCCC1",
"CC1CCCCC1",
"c1ccoc1",
]
tr, te = scaffold_split(smi, 3, 0.7)
scaff = lambda inds: {
MurckoScaffold.MurckoScaffoldSmiles(smiles=smi[i]) for i in inds
}
assert not scaff(tr) & scaff(te)
def test_catalog_conflicting_identifiers(tmp_path):
f = tmp_path / "parents.csv"
f.write_text("id,smiles\na,CC\na,CCC\n")
with pytest.raises(ValueError, match="Conflicting"):
read_catalog(f)
def test_resumed_training_matches_uninterrupted(tmp_path):
g = toy_graph(3)
rewards = {"A": 0.0, "B": 0.0}
full, _ = train(g, rewards, steps=30, seed=2, batch_size=16)
train(g, rewards, steps=12, seed=2, batch_size=16, output=tmp_path)
resumed, _ = train(g, rewards, steps=30, seed=2, batch_size=16, resume=tmp_path)
np.testing.assert_allclose(full.probabilities(), resumed.probabilities(), atol=1e-9)
def test_replay_requires_parent_and_stop():
assert replay(
{
"actions": [{"kind": "parent", "smiles": "CC"}, {"kind": "stop"}],
"outcome": "CC",
},
0,
)
assert not replay(
{"actions": [{"kind": "parent", "smiles": "CC"}], "outcome": "CC"}, 0
)
assert not replay(
{
"actions": [
{"kind": "parent", "smiles": "CC"},
{"kind": "parent", "smiles": "CC"},
{"kind": "stop"},
],
"outcome": "CC",
},
2,
)
def test_terminal_merging_ablation_changes_endpoint_mass():
from dooable.ablations import duplicate_endpoint_policy
g = toy_graph(5)
r = {"A": 0.0, "B": 0.0}
p = endpoint_distribution(g, duplicate_endpoint_policy(g, r, 0.7))
assert p["A"] == pytest.approx(5 / 6)
|