synomega 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (44) hide show
  1. synomega-0.1.0/LICENSE +21 -0
  2. synomega-0.1.0/PKG-INFO +168 -0
  3. synomega-0.1.0/README.md +138 -0
  4. synomega-0.1.0/pyproject.toml +48 -0
  5. synomega-0.1.0/setup.cfg +4 -0
  6. synomega-0.1.0/src/synomega/__init__.py +55 -0
  7. synomega-0.1.0/src/synomega/chem/__init__.py +20 -0
  8. synomega-0.1.0/src/synomega/chem/features.py +150 -0
  9. synomega-0.1.0/src/synomega/chem/mol.py +142 -0
  10. synomega-0.1.0/src/synomega/chem/reaction.py +99 -0
  11. synomega-0.1.0/src/synomega/chem/template.py +159 -0
  12. synomega-0.1.0/src/synomega/cli.py +144 -0
  13. synomega-0.1.0/src/synomega/planner.py +97 -0
  14. synomega-0.1.0/src/synomega/route/__init__.py +5 -0
  15. synomega-0.1.0/src/synomega/route/route.py +258 -0
  16. synomega-0.1.0/src/synomega/search/__init__.py +50 -0
  17. synomega-0.1.0/src/synomega/search/base.py +150 -0
  18. synomega-0.1.0/src/synomega/search/bfs.py +139 -0
  19. synomega-0.1.0/src/synomega/search/graph.py +208 -0
  20. synomega-0.1.0/src/synomega/search/mcts.py +267 -0
  21. synomega-0.1.0/src/synomega/search/retrostar.py +212 -0
  22. synomega-0.1.0/src/synomega/search/value.py +100 -0
  23. synomega-0.1.0/src/synomega/singlestep/__init__.py +27 -0
  24. synomega-0.1.0/src/synomega/singlestep/_dmpnn.py +114 -0
  25. synomega-0.1.0/src/synomega/singlestep/base.py +57 -0
  26. synomega-0.1.0/src/synomega/singlestep/cache.py +169 -0
  27. synomega-0.1.0/src/synomega/singlestep/template_gnn.py +280 -0
  28. synomega-0.1.0/src/synomega/singlestep/template_rule.py +70 -0
  29. synomega-0.1.0/src/synomega/stock/__init__.py +7 -0
  30. synomega-0.1.0/src/synomega/stock/base.py +55 -0
  31. synomega-0.1.0/src/synomega/stock/inmemory.py +130 -0
  32. synomega-0.1.0/src/synomega/stock/sqlite.py +118 -0
  33. synomega-0.1.0/src/synomega/synthesizability/__init__.py +6 -0
  34. synomega-0.1.0/src/synomega/synthesizability/metrics.py +204 -0
  35. synomega-0.1.0/src/synomega/synthesizability/scorer.py +176 -0
  36. synomega-0.1.0/src/synomega.egg-info/PKG-INFO +168 -0
  37. synomega-0.1.0/src/synomega.egg-info/SOURCES.txt +42 -0
  38. synomega-0.1.0/src/synomega.egg-info/dependency_links.txt +1 -0
  39. synomega-0.1.0/src/synomega.egg-info/entry_points.txt +2 -0
  40. synomega-0.1.0/src/synomega.egg-info/requires.txt +13 -0
  41. synomega-0.1.0/src/synomega.egg-info/top_level.txt +1 -0
  42. synomega-0.1.0/tests/test_chem.py +67 -0
  43. synomega-0.1.0/tests/test_search.py +122 -0
  44. synomega-0.1.0/tests/test_synthesizability.py +153 -0
synomega-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 zbc0315
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,168 @@
1
+ Metadata-Version: 2.4
2
+ Name: synomega
3
+ Version: 0.1.0
4
+ Summary: Retrosynthesis toolkit: single-step prediction, multi-step route planning, synthesizability scoring
5
+ Author: zbc0315
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/zbc0315/synomega
8
+ Project-URL: Repository, https://github.com/zbc0315/synomega
9
+ Project-URL: Issues, https://github.com/zbc0315/synomega/issues
10
+ Keywords: retrosynthesis,cheminformatics,synthesizability,rdkit,route-planning
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: License :: OSI Approved :: MIT License
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: Topic :: Scientific/Engineering :: Chemistry
15
+ Classifier: Operating System :: OS Independent
16
+ Requires-Python: >=3.10
17
+ Description-Content-Type: text/markdown
18
+ License-File: LICENSE
19
+ Requires-Dist: rdkit>=2023.3
20
+ Requires-Dist: numpy>=1.23
21
+ Provides-Extra: gnn
22
+ Requires-Dist: torch>=2.0; extra == "gnn"
23
+ Requires-Dist: torch_geometric>=2.4; extra == "gnn"
24
+ Requires-Dist: pyyaml>=6.0; extra == "gnn"
25
+ Provides-Extra: viz
26
+ Requires-Dist: graphviz>=0.20; extra == "viz"
27
+ Provides-Extra: dev
28
+ Requires-Dist: pytest>=7.0; extra == "dev"
29
+ Dynamic: license-file
30
+
31
+ # synomega
32
+
33
+ Retrosynthesis toolkit: **single-step prediction → multi-step route planning → synthesizability scoring**.
34
+
35
+ ```
36
+ synthesizability is this target reachable from purchasable material, in N steps?
37
+
38
+ search Retro* / MCTS / best-first over an AND-OR graph
39
+
40
+ singlestep product SMILES -> ranked reactant candidates
41
+ ```
42
+
43
+ The layers are decoupled by a deliberately narrow interface: a single-step
44
+ backend only implements `predict(smiles, top_k) -> [Prediction]`. Whether it is
45
+ a graph neural network, a transformer, or plain template matching is invisible
46
+ to the planner.
47
+
48
+ ## Install
49
+
50
+ ```bash
51
+ pip install -e . # core: rdkit + numpy, no torch
52
+ pip install -e '.[gnn]' # adds the D-MPNN neural backend
53
+ ```
54
+
55
+ The neural backend is an **optional** extra on purpose — the template-rule
56
+ backend runs anywhere, with no GPU and no torch.
57
+
58
+ ## Quick start
59
+
60
+ ```python
61
+ from synomega import Planner, SynthesizabilityScorer
62
+ from synomega.singlestep import TemplateGNN
63
+ from synomega.stock import InMemoryStock
64
+
65
+ model = TemplateGNN.from_pretrained("../ml-template-gnn/runs/uspto50k_r0_min10")
66
+ stock = InMemoryStock.from_keys_file("emolecules.keys.gz")
67
+ planner = Planner(model, stock, algorithm="retrostar")
68
+
69
+ result = planner.plan("CC(=O)Nc1ccccc1", max_depth=5, time_limit=60)
70
+ print(result.solved)
71
+ print(result.best_route.describe())
72
+ ```
73
+
74
+ ```
75
+ target: CC(=O)Nc1ccccc1
76
+ solved: True steps: 2 depth: 2 bb_coverage: 1.00
77
+ [1] CC(=O)O.Nc1ccccc1>>CC(=O)Nc1ccccc1 (score=0.4348)
78
+ [2] O=[N+]([O-])c1ccccc1>>Nc1ccccc1 (score=0.2174)
79
+ ```
80
+
81
+ ### Synthesizability
82
+
83
+ ```python
84
+ scorer = SynthesizabilityScorer(planner)
85
+
86
+ r = scorer.score("CC(=O)Nc1ccccc1", max_steps=5)
87
+ r.solved # True — a complete route to purchasable material exists
88
+ r.bb_coverage # 1.0 — fraction of leaves that are buyable
89
+ r.min_depth # 2 — steps in the shortest solved route
90
+
91
+ report = scorer.score_batch(targets, max_steps=5)
92
+ report.solve_rate # headline benchmark number
93
+ report.mean_bb_coverage
94
+ report.to_dataframe()
95
+ ```
96
+
97
+ ## The two synthesizability metrics
98
+
99
+ These get conflated in the literature; synomega keeps them apart because they
100
+ answer different questions.
101
+
102
+ | Metric | Meaning | Use it for |
103
+ |---|---|---|
104
+ | `solved@N` / `solve_rate` | Binary: does a route of depth ≤ N exist whose leaves are **all** purchasable? | Comparing against published numbers |
105
+ | `bb_coverage@N` | Continuous: fraction of the best route's leaves that are purchasable | Ranking molecules by how close they are |
106
+
107
+ `bb_coverage` matters because most targets are unsolved at realistic step
108
+ limits. A 5-step route with 4 of 5 leaves buyable scores 0.8, not 0 — so a
109
+ near-miss is distinguishable from a total failure.
110
+
111
+ ## CLI
112
+
113
+ ```bash
114
+ # one-time: precompute InChIKeys so later loads take seconds, not minutes
115
+ synomega build-stock --catalogue emolecules.smi.gz --out emolecules.keys.gz
116
+
117
+ synomega plan --target "CC(=O)Nc1ccccc1" --model runs/uspto50k_r0_min10 \
118
+ --stock emolecules.keys.gz --stock-is-keys --max-steps 5
119
+
120
+ synomega score --targets targets.smi --model runs/uspto50k_r0_min10 \
121
+ --stock emolecules.keys.gz --stock-is-keys \
122
+ --max-steps 5 --out report.json
123
+ ```
124
+
125
+ ## Search algorithms
126
+
127
+ | Algorithm | Character | When |
128
+ |---|---|---|
129
+ | `retrostar` | Expands the frontier molecule with the lowest estimated **total route cost** (Chen et al. 2020) | Default |
130
+ | `mcts` | UCT with greedy rollouts; tolerates an unreliable top-1 | When the single-step model is weak |
131
+ | `bfs` | Best-first on `g + h` | Baseline, debugging |
132
+
133
+ All three share the AND-OR graph, the budget, and the route extractor, so they
134
+ are directly comparable.
135
+
136
+ ## Design notes
137
+
138
+ **AND-OR graph, not a tree.** A molecule node is solved if it is in stock or
139
+ *any* of its reactions is solved; a reaction is solved if *all* its reactants
140
+ are. Molecules are interned by InChIKey, so an intermediate reached down two
141
+ branches is one node expanded once. Cycles are rejected when the edge is
142
+ created — a route that makes X from X is not a route.
143
+
144
+ **Batched expansion.** Search is naturally "expand one node, call the model
145
+ once", which leaves a GPU idle. The algorithms pull a batch of frontier nodes
146
+ and issue one `predict_batch`.
147
+
148
+ **Caching.** Search revisits the same molecule constantly. `Planner(cache=True)`
149
+ (the default) memoizes expansions; `cache_path=` persists them to SQLite so the
150
+ cache survives across runs.
151
+
152
+ **Stock membership is by InChIKey**, so keys computed here match a vendor
153
+ catalogue written by a different toolkit. There is deliberately no Bloom-filter
154
+ backend: false positives would inflate solve-rate and quietly break
155
+ comparability with published numbers.
156
+
157
+ ## Status
158
+
159
+ Implemented and tested: chem layer, single-step interface + template-rule and
160
+ D-MPNN backends, expansion cache, AND-OR graph, all three search algorithms,
161
+ route extraction/scoring/serialization, synthesizability metrics, CLI.
162
+
163
+ Reserved but not implemented: reaction conditions (`RouteStep.conditions` is
164
+ always `None`); a condition model can fill it without touching the search layer.
165
+
166
+ ```bash
167
+ pytest # 43 tests
168
+ ```
@@ -0,0 +1,138 @@
1
+ # synomega
2
+
3
+ Retrosynthesis toolkit: **single-step prediction → multi-step route planning → synthesizability scoring**.
4
+
5
+ ```
6
+ synthesizability is this target reachable from purchasable material, in N steps?
7
+
8
+ search Retro* / MCTS / best-first over an AND-OR graph
9
+
10
+ singlestep product SMILES -> ranked reactant candidates
11
+ ```
12
+
13
+ The layers are decoupled by a deliberately narrow interface: a single-step
14
+ backend only implements `predict(smiles, top_k) -> [Prediction]`. Whether it is
15
+ a graph neural network, a transformer, or plain template matching is invisible
16
+ to the planner.
17
+
18
+ ## Install
19
+
20
+ ```bash
21
+ pip install -e . # core: rdkit + numpy, no torch
22
+ pip install -e '.[gnn]' # adds the D-MPNN neural backend
23
+ ```
24
+
25
+ The neural backend is an **optional** extra on purpose — the template-rule
26
+ backend runs anywhere, with no GPU and no torch.
27
+
28
+ ## Quick start
29
+
30
+ ```python
31
+ from synomega import Planner, SynthesizabilityScorer
32
+ from synomega.singlestep import TemplateGNN
33
+ from synomega.stock import InMemoryStock
34
+
35
+ model = TemplateGNN.from_pretrained("../ml-template-gnn/runs/uspto50k_r0_min10")
36
+ stock = InMemoryStock.from_keys_file("emolecules.keys.gz")
37
+ planner = Planner(model, stock, algorithm="retrostar")
38
+
39
+ result = planner.plan("CC(=O)Nc1ccccc1", max_depth=5, time_limit=60)
40
+ print(result.solved)
41
+ print(result.best_route.describe())
42
+ ```
43
+
44
+ ```
45
+ target: CC(=O)Nc1ccccc1
46
+ solved: True steps: 2 depth: 2 bb_coverage: 1.00
47
+ [1] CC(=O)O.Nc1ccccc1>>CC(=O)Nc1ccccc1 (score=0.4348)
48
+ [2] O=[N+]([O-])c1ccccc1>>Nc1ccccc1 (score=0.2174)
49
+ ```
50
+
51
+ ### Synthesizability
52
+
53
+ ```python
54
+ scorer = SynthesizabilityScorer(planner)
55
+
56
+ r = scorer.score("CC(=O)Nc1ccccc1", max_steps=5)
57
+ r.solved # True — a complete route to purchasable material exists
58
+ r.bb_coverage # 1.0 — fraction of leaves that are buyable
59
+ r.min_depth # 2 — steps in the shortest solved route
60
+
61
+ report = scorer.score_batch(targets, max_steps=5)
62
+ report.solve_rate # headline benchmark number
63
+ report.mean_bb_coverage
64
+ report.to_dataframe()
65
+ ```
66
+
67
+ ## The two synthesizability metrics
68
+
69
+ These get conflated in the literature; synomega keeps them apart because they
70
+ answer different questions.
71
+
72
+ | Metric | Meaning | Use it for |
73
+ |---|---|---|
74
+ | `solved@N` / `solve_rate` | Binary: does a route of depth ≤ N exist whose leaves are **all** purchasable? | Comparing against published numbers |
75
+ | `bb_coverage@N` | Continuous: fraction of the best route's leaves that are purchasable | Ranking molecules by how close they are |
76
+
77
+ `bb_coverage` matters because most targets are unsolved at realistic step
78
+ limits. A 5-step route with 4 of 5 leaves buyable scores 0.8, not 0 — so a
79
+ near-miss is distinguishable from a total failure.
80
+
81
+ ## CLI
82
+
83
+ ```bash
84
+ # one-time: precompute InChIKeys so later loads take seconds, not minutes
85
+ synomega build-stock --catalogue emolecules.smi.gz --out emolecules.keys.gz
86
+
87
+ synomega plan --target "CC(=O)Nc1ccccc1" --model runs/uspto50k_r0_min10 \
88
+ --stock emolecules.keys.gz --stock-is-keys --max-steps 5
89
+
90
+ synomega score --targets targets.smi --model runs/uspto50k_r0_min10 \
91
+ --stock emolecules.keys.gz --stock-is-keys \
92
+ --max-steps 5 --out report.json
93
+ ```
94
+
95
+ ## Search algorithms
96
+
97
+ | Algorithm | Character | When |
98
+ |---|---|---|
99
+ | `retrostar` | Expands the frontier molecule with the lowest estimated **total route cost** (Chen et al. 2020) | Default |
100
+ | `mcts` | UCT with greedy rollouts; tolerates an unreliable top-1 | When the single-step model is weak |
101
+ | `bfs` | Best-first on `g + h` | Baseline, debugging |
102
+
103
+ All three share the AND-OR graph, the budget, and the route extractor, so they
104
+ are directly comparable.
105
+
106
+ ## Design notes
107
+
108
+ **AND-OR graph, not a tree.** A molecule node is solved if it is in stock or
109
+ *any* of its reactions is solved; a reaction is solved if *all* its reactants
110
+ are. Molecules are interned by InChIKey, so an intermediate reached down two
111
+ branches is one node expanded once. Cycles are rejected when the edge is
112
+ created — a route that makes X from X is not a route.
113
+
114
+ **Batched expansion.** Search is naturally "expand one node, call the model
115
+ once", which leaves a GPU idle. The algorithms pull a batch of frontier nodes
116
+ and issue one `predict_batch`.
117
+
118
+ **Caching.** Search revisits the same molecule constantly. `Planner(cache=True)`
119
+ (the default) memoizes expansions; `cache_path=` persists them to SQLite so the
120
+ cache survives across runs.
121
+
122
+ **Stock membership is by InChIKey**, so keys computed here match a vendor
123
+ catalogue written by a different toolkit. There is deliberately no Bloom-filter
124
+ backend: false positives would inflate solve-rate and quietly break
125
+ comparability with published numbers.
126
+
127
+ ## Status
128
+
129
+ Implemented and tested: chem layer, single-step interface + template-rule and
130
+ D-MPNN backends, expansion cache, AND-OR graph, all three search algorithms,
131
+ route extraction/scoring/serialization, synthesizability metrics, CLI.
132
+
133
+ Reserved but not implemented: reaction conditions (`RouteStep.conditions` is
134
+ always `None`); a condition model can fill it without touching the search layer.
135
+
136
+ ```bash
137
+ pytest # 43 tests
138
+ ```
@@ -0,0 +1,48 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "synomega"
7
+ version = "0.1.0"
8
+ description = "Retrosynthesis toolkit: single-step prediction, multi-step route planning, synthesizability scoring"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = { text = "MIT" }
12
+ authors = [{ name = "zbc0315" }]
13
+ keywords = ["retrosynthesis", "cheminformatics", "synthesizability", "rdkit", "route-planning"]
14
+ classifiers = [
15
+ "Programming Language :: Python :: 3",
16
+ "License :: OSI Approved :: MIT License",
17
+ "Intended Audience :: Science/Research",
18
+ "Topic :: Scientific/Engineering :: Chemistry",
19
+ "Operating System :: OS Independent",
20
+ ]
21
+ dependencies = [
22
+ "rdkit>=2023.3",
23
+ "numpy>=1.23",
24
+ ]
25
+
26
+ [project.urls]
27
+ Homepage = "https://github.com/zbc0315/synomega"
28
+ Repository = "https://github.com/zbc0315/synomega"
29
+ Issues = "https://github.com/zbc0315/synomega/issues"
30
+
31
+ [project.optional-dependencies]
32
+ # Neural single-step backends (D-MPNN template classifier).
33
+ gnn = [
34
+ "torch>=2.0",
35
+ "torch_geometric>=2.4",
36
+ "pyyaml>=6.0",
37
+ ]
38
+ viz = ["graphviz>=0.20"]
39
+ dev = ["pytest>=7.0"]
40
+
41
+ [project.scripts]
42
+ synomega = "synomega.cli:main"
43
+
44
+ [tool.setuptools.packages.find]
45
+ where = ["src"]
46
+
47
+ [tool.pytest.ini_options]
48
+ testpaths = ["tests"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,55 @@
1
+ """synomega — retrosynthesis toolkit.
2
+
3
+ Three layers, deliberately decoupled:
4
+
5
+ synthesizability is this target reachable from purchasable material?
6
+
7
+ search multi-step route planning over an AND-OR graph
8
+
9
+ singlestep product SMILES -> ranked reactant candidates
10
+
11
+ Quick start::
12
+
13
+ from synomega import Planner, SynthesizabilityScorer
14
+ from synomega.singlestep import TemplateGNN
15
+ from synomega.stock import InMemoryStock
16
+
17
+ model = TemplateGNN.from_pretrained("runs/uspto50k_r0_min10")
18
+ stock = InMemoryStock.from_file("emolecules.smi")
19
+ planner = Planner(model, stock, algorithm="retrostar")
20
+
21
+ result = planner.plan("CC(=O)Nc1ccccc1", max_depth=5)
22
+ print(result.best_route.describe())
23
+
24
+ scorer = SynthesizabilityScorer(planner)
25
+ print(scorer.score("CC(=O)Nc1ccccc1", max_steps=5))
26
+ """
27
+
28
+ from __future__ import annotations
29
+
30
+ __version__ = "0.1.0"
31
+
32
+ from .chem import Molecule, Reaction
33
+ from .planner import Planner
34
+ from .route import Route
35
+ from .search import Budget, SearchResult
36
+ from .singlestep import Prediction, SingleStepModel
37
+ from .stock import BuildingBlockSet, InMemoryStock
38
+ from .synthesizability import BatchReport, MoleculeReport, SynthesizabilityScorer
39
+
40
+ __all__ = [
41
+ "__version__",
42
+ "Planner",
43
+ "SynthesizabilityScorer",
44
+ "MoleculeReport",
45
+ "BatchReport",
46
+ "Route",
47
+ "SearchResult",
48
+ "Budget",
49
+ "Molecule",
50
+ "Reaction",
51
+ "Prediction",
52
+ "SingleStepModel",
53
+ "BuildingBlockSet",
54
+ "InMemoryStock",
55
+ ]
@@ -0,0 +1,20 @@
1
+ """Chemistry primitives: molecules, reactions, retro-template application."""
2
+
3
+ from .mol import Molecule, MoleculeError, canonicalize, inchi_key, split_components
4
+ from .reaction import Conditions, Reaction, extract_product, parse_reaction_smiles
5
+ from .template import TemplateLibrary, TemplateOutcome, apply_template
6
+
7
+ __all__ = [
8
+ "Molecule",
9
+ "MoleculeError",
10
+ "canonicalize",
11
+ "inchi_key",
12
+ "split_components",
13
+ "Reaction",
14
+ "Conditions",
15
+ "parse_reaction_smiles",
16
+ "extract_product",
17
+ "TemplateLibrary",
18
+ "TemplateOutcome",
19
+ "apply_template",
20
+ ]
@@ -0,0 +1,150 @@
1
+ """Atom/bond featurization for the D-MPNN template classifier.
2
+
3
+ ⚠️ This is a VERBATIM mirror of `ml-template-gnn/src/template_gnn/featurize.py`.
4
+ It is vendored so synomega can run inference from a checkpoint without depending
5
+ on the training repo. It MUST stay bit-identical to whatever produced the
6
+ checkpoint you load — any drift silently corrupts predictions rather than
7
+ raising, so `synomega.singlestep.template_gnn` asserts ATOM_FDIM against the
8
+ checkpoint's first-layer weight shape.
9
+
10
+ Atom features (45 dim), bond features (12 dim). See the training repo for the
11
+ full field-by-field breakdown.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import math
17
+
18
+ from rdkit import Chem
19
+ from rdkit.Chem import AllChem, rdchem
20
+
21
+ ATOM_LIST = [6, 7, 8, 9, 15, 16, 17, 35, 53]
22
+ DEGREE_LIST = list(range(6))
23
+ CHARGE_LIST = [-2, -1, 0, 1, 2]
24
+ CHIRALITY_LIST = [
25
+ rdchem.ChiralType.CHI_UNSPECIFIED,
26
+ rdchem.ChiralType.CHI_TETRAHEDRAL_CW,
27
+ rdchem.ChiralType.CHI_TETRAHEDRAL_CCW,
28
+ rdchem.ChiralType.CHI_OTHER,
29
+ ]
30
+ NUM_H_LIST = list(range(5))
31
+ HYBRIDIZATION_LIST = [
32
+ rdchem.HybridizationType.SP,
33
+ rdchem.HybridizationType.SP2,
34
+ rdchem.HybridizationType.SP3,
35
+ rdchem.HybridizationType.SP3D,
36
+ rdchem.HybridizationType.SP3D2,
37
+ ]
38
+
39
+ ATOM_FDIM = (
40
+ len(ATOM_LIST) + 1
41
+ + len(DEGREE_LIST) + 1
42
+ + len(CHARGE_LIST) + 1
43
+ + len(CHIRALITY_LIST)
44
+ + len(NUM_H_LIST) + 1
45
+ + len(HYBRIDIZATION_LIST) + 1
46
+ + 1 # is_aromatic
47
+ + 1 # is_in_ring
48
+ + 1 # mass
49
+ + 1 # chi_pauling / 4
50
+ + 1 # q_gasteiger (masked)
51
+ + 1 # q_valid
52
+ )
53
+
54
+ PAULING_EN = {
55
+ 1: 2.20, 3: 0.98, 4: 1.57, 5: 2.04, 6: 2.55, 7: 3.04, 8: 3.44, 9: 3.98,
56
+ 11: 0.93, 12: 1.31, 13: 1.61, 14: 1.90, 15: 2.19, 16: 2.58, 17: 3.16,
57
+ 19: 0.82, 20: 1.00, 26: 1.83, 27: 1.88, 28: 1.91, 29: 1.90, 30: 1.65,
58
+ 33: 2.18, 34: 2.55, 35: 2.96, 47: 1.93, 50: 1.96, 53: 2.66,
59
+ 78: 2.28, 79: 2.54, 80: 2.00,
60
+ }
61
+
62
+ BOND_TYPE_LIST = [
63
+ rdchem.BondType.SINGLE,
64
+ rdchem.BondType.DOUBLE,
65
+ rdchem.BondType.TRIPLE,
66
+ rdchem.BondType.AROMATIC,
67
+ ]
68
+ STEREO_LIST = [
69
+ rdchem.BondStereo.STEREONONE,
70
+ rdchem.BondStereo.STEREOANY,
71
+ rdchem.BondStereo.STEREOZ,
72
+ rdchem.BondStereo.STEREOE,
73
+ rdchem.BondStereo.STEREOCIS,
74
+ rdchem.BondStereo.STEREOTRANS,
75
+ ]
76
+
77
+ BOND_FDIM = len(BOND_TYPE_LIST) + 1 + 1 + len(STEREO_LIST)
78
+
79
+
80
+ def _one_hot(value, choices, allow_other: bool = True) -> list[int]:
81
+ feat = [0] * (len(choices) + (1 if allow_other else 0))
82
+ try:
83
+ idx = choices.index(value)
84
+ except ValueError:
85
+ if not allow_other:
86
+ return feat
87
+ idx = len(choices)
88
+ feat[idx] = 1
89
+ return feat
90
+
91
+
92
+ def compute_gasteiger(mol: Chem.Mol) -> bool:
93
+ """Annotate Gasteiger charges in place; False when RDKit refuses."""
94
+ try:
95
+ AllChem.ComputeGasteigerCharges(mol)
96
+ return True
97
+ except Exception:
98
+ return False
99
+
100
+
101
+ def atom_features(atom: Chem.Atom) -> list[float]:
102
+ z = atom.GetAtomicNum()
103
+ chi = PAULING_EN.get(z, 2.0) / 4.0
104
+
105
+ q = 0.0
106
+ q_valid = 0
107
+ if atom.HasProp("_GasteigerCharge"):
108
+ raw = atom.GetDoubleProp("_GasteigerCharge")
109
+ if math.isfinite(raw):
110
+ q = raw
111
+ q_valid = 1
112
+
113
+ return (
114
+ _one_hot(z, ATOM_LIST)
115
+ + _one_hot(atom.GetTotalDegree(), DEGREE_LIST)
116
+ + _one_hot(atom.GetFormalCharge(), CHARGE_LIST)
117
+ + _one_hot(atom.GetChiralTag(), CHIRALITY_LIST, allow_other=False)
118
+ + _one_hot(atom.GetTotalNumHs(), NUM_H_LIST)
119
+ + _one_hot(atom.GetHybridization(), HYBRIDIZATION_LIST)
120
+ + [1 if atom.GetIsAromatic() else 0]
121
+ + [1 if atom.IsInRing() else 0]
122
+ + [atom.GetMass() * 0.01]
123
+ + [chi, q, float(q_valid)]
124
+ )
125
+
126
+
127
+ def bond_features(bond: Chem.Bond) -> list[float]:
128
+ return (
129
+ _one_hot(bond.GetBondType(), BOND_TYPE_LIST, allow_other=False)
130
+ + [1 if bond.GetIsConjugated() else 0]
131
+ + [1 if bond.IsInRing() else 0]
132
+ + _one_hot(bond.GetStereo(), STEREO_LIST, allow_other=False)
133
+ )
134
+
135
+
136
+ def largest_fragment(mol: Chem.Mol) -> Chem.Mol:
137
+ frags = Chem.GetMolFrags(mol, asMols=True, sanitizeFrags=False)
138
+ if len(frags) == 1:
139
+ return mol
140
+ return max(frags, key=lambda m: m.GetNumHeavyAtoms())
141
+
142
+
143
+ __all__ = [
144
+ "ATOM_FDIM",
145
+ "BOND_FDIM",
146
+ "atom_features",
147
+ "bond_features",
148
+ "compute_gasteiger",
149
+ "largest_fragment",
150
+ ]