-
Notifications
You must be signed in to change notification settings - Fork 42
/
moler_decoding_utils.py
748 lines (643 loc) · 30.2 KB
/
moler_decoding_utils.py
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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
"""Decoding utilities for a MoLeR."""
from collections import defaultdict
from dataclasses import dataclass
from enum import Enum
from typing import Any, List, Optional, Tuple, NamedTuple, Iterable, DefaultDict, Union
import heapq
import numpy as np
import tensorflow as tf
from rdkit import Chem
from molecule_generation.chem.atom_feature_utils import AtomFeatureExtractor
from molecule_generation.chem.molecule_dataset_utils import BOND_DICT, featurise_atoms
from molecule_generation.chem.rdkit_helpers import initialise_atom_from_symbol, get_atom_symbol
from molecule_generation.chem.motif_utils import (
MotifAnnotation,
MotifAtomAnnotation,
MotifVocabulary,
)
from molecule_generation.chem.topology_features import calculate_topology_features
from molecule_generation.chem.valence_constraints import (
constrain_edge_choices_based_on_valence,
constrain_edge_types_based_on_valence,
)
from molecule_generation.preprocessing.cgvae_generation_trace import (
calculate_dist_from_focus_to_valid_target,
)
from molecule_generation.preprocessing.moler_generation_trace import (
Edge,
get_open_attachment_points,
)
EDGE_TYPE_IDX_TO_BOND_TYPE = {
0: Chem.BondType.SINGLE,
1: Chem.BondType.DOUBLE,
2: Chem.BondType.TRIPLE,
}
class DecoderSamplingMode(Enum):
GREEDY = 0
SAMPLING = 1
def sample_indices_from_logprobs(
num_samples: int, sampling_mode: DecoderSamplingMode, logprobs: np.ndarray
) -> Iterable[int]:
"""Samples indices (without replacement) given the log-likelihoods.
Args:
num_samples: intended number of samples
sampling_mode: sampling method (greedy
logprobs: log-probabilities of selecting appropriate entries
Returns:
indices of picked values, shape (n,), where n = min(num_samples, available_samples)
Note:
the ordering of returned indices is arbitrary
"""
num_choices = logprobs.shape[0]
indices = np.arange(num_choices)
num_samples = min(num_samples, num_choices) # Handle cases where we only have few candidates
if sampling_mode == DecoderSamplingMode.GREEDY:
# Note that this will return the top num_samples indices, but not in order:
picked_indices = np.argpartition(logprobs, -num_samples)[-num_samples:]
elif sampling_mode == DecoderSamplingMode.SAMPLING:
p = np.exp(logprobs) # Convert to probabilities
# We can only sample values with non-zero probabilities
num_choices = np.sum(p > 0)
num_samples = min(num_samples, num_choices)
picked_indices = np.random.choice(
indices,
size=(num_samples,),
replace=False,
p=p,
)
else:
raise ValueError(f"Sampling method {sampling_mode} not known.")
return picked_indices
class MoleculeGenerationEdgeCandidateInfo(NamedTuple):
target_node_idx: int
score: float
logprob: float
correct: Optional[bool]
type_idx_to_logprobs: np.ndarray
@dataclass
class MoleculeGenerationAtomChoiceInfo:
node_idx: int
true_type_idx: Optional[List[int]]
type_idx_to_prob: List[float]
molecule: Optional[Chem.Mol] = None
@dataclass
class MoleculeGenerationEdgeChoiceInfo:
focus_node_idx: int
partial_molecule_adjacency_lists: List[np.ndarray]
candidate_edge_infos: List[MoleculeGenerationEdgeCandidateInfo]
no_edge_score: float
no_edge_logprob: float
no_edge_correct: Optional[bool]
molecule: Optional[Chem.Mol] = None
@dataclass
class MoleculeGenerationAttachmentPointChoiceInfo:
partial_molecule_adjacency_lists: List[np.ndarray]
motif_nodes: List[int]
candidate_attachment_points: List[int]
candidate_idx_to_prob: List[float]
correct_attachment_point_idx: Optional[int]
molecule: Optional[Chem.Mol] = None
MoleculeGenerationChoiceInfo = Union[
MoleculeGenerationAtomChoiceInfo,
MoleculeGenerationEdgeChoiceInfo,
MoleculeGenerationAttachmentPointChoiceInfo,
]
class MoLeRDecoderState(object):
"""
Data class holding the state of the MoLeR decoder.
Note: To enable safe usage in contexts with sharing (such as beam search), the
contract is that mutable members of this class cannot be manipulated directly,
but need to be copied on modification. Hence, a number of static creation
methods (new_with_*) are provided which should enable all necessary actions.
"""
def __init__(
self,
molecule_representation: tf.Tensor,
molecule_id: Any,
logprob: float = 0.0,
molecule: Optional[Chem.Mol] = None,
atom_types: Optional[List[str]] = None,
adjacency_lists: Optional[List[List[Tuple[int, int]]]] = None,
visited_atoms: Optional[List[int]] = None,
atoms_to_visit: Optional[List[int]] = None,
atoms_to_mark_as_visited: Optional[List[int]] = None,
focus_atom: Optional[int] = None,
prior_focus_atom: Optional[int] = None,
generation_steps: Optional[List[MoleculeGenerationChoiceInfo]] = None,
candidate_attachment_points: Optional[List[int]] = None,
motifs: Optional[List[MotifAnnotation]] = None,
num_free_bond_slots: Optional[List[Optional[int]]] = None,
):
"""
Create new state tracking generation of a molecule.
Args:
molecule_representation: Representation from which this molecule
is being generated.
molecule_id: Object used to identify this particular input; usually
set at the beginning of the decoding process and then simply
propagated.
logprob: Log probability of the current decoder state, i.e., the sum
of the log probabilities of all steps leading to it from the
initial start.
molecule: An optional (partial) molecule encoding the current
state of the generation. If it is missing, we start from an
empty molecule, and all remaining arguments have to be None
as well.
atom_types: A list of the atom types (string representation including
charges, e.g., "N++") used in the passed partial molecule. This has
to be passed in if the given partial molecule is non-empty.
adjacency_lists: _Symmetric_ adjacency lists (i.e., we assume
that both (u, v) and (v, u) are provided), corresponding to the
partial molecule we are decoding, as a list of int32 tensors of
shape [E, 2].
visited_atoms: Atoms in the molecule (stored as indices) that we have
already processed and to which we can add fresh connections.
atoms_to_visit: Atoms in the molecule (stored as indices) that we still
need to process and from which we can add new bonds.
atoms_to_mark_as_visited: Atoms in the molecule (stored as indices) that we want to mark
as visited when `new_with_focus_marked_as_visited` is called. This is used to delay
marking as visited the motif nodes, in order to avoid adding additional edges within
motifs.
focus_atom: Atom that we are currently focusing on, i.e., from which
we generate new bonds.
prior_focus_atom: The focus atom that we finished last. This is required
to inform the next-atom choice...
generation_steps: List of generation step metadata, or `None` if we don't want to store
the generation traces.
candidate_attachment_points: Atoms in the molecule (stored as indices) that are
candidates to consider for picking an attachment point.
motifs: List of annotations of motifs that already exist in the current molecule.
num_free_bond_slots: Number of bonds that can still be attached to each existing atom.
Passing `None` for a specific atom means there are no constraints for it, passing
`None` for `num_free_bond_slots` itself means there are no constraints at all.
"""
self._molecule_representation = molecule_representation
self._molecule_id = molecule_id
self._logprob = logprob
if molecule is None and not (
atom_types is None
and adjacency_lists is None
and visited_atoms is None
and atoms_to_visit is None
and focus_atom is None
and prior_focus_atom is None
):
raise ValueError(
"Molecule generation states created without a partial molecule cannot take any other values."
)
self._molecule = molecule or Chem.RWMol()
if self._molecule.GetNumAtoms() > 0 and atom_types is None:
raise ValueError(
f"Molecule generation states starting with a partial molecule require atom type information!"
)
self._atom_types: List[str] = atom_types or []
self._adjacency_lists = adjacency_lists or [[] for _ in range(len(BOND_DICT))]
self._visited_atoms: List[int] = visited_atoms or []
self._atoms_to_visit: List[int] = atoms_to_visit or []
self._atoms_to_mark_as_visited: List[int] = atoms_to_mark_as_visited or []
self._focus_atom: Optional[int] = focus_atom
self._prior_focus_atom: Optional[int] = prior_focus_atom
self._generation_steps = generation_steps
self._candidate_attachment_points = candidate_attachment_points or []
self._motifs = motifs or []
self._num_free_bond_slots = num_free_bond_slots or [None] * len(atom_types)
@staticmethod
def extend_generation_steps(
generation_steps: Optional[List[MoleculeGenerationChoiceInfo]],
choice_info: MoleculeGenerationChoiceInfo,
molecule: Chem.Mol,
):
if choice_info is None:
return generation_steps
else:
assert generation_steps is not None
choice_info.molecule = Chem.Mol(molecule)
return list(generation_steps) + [choice_info]
@staticmethod
def new_with_added_atom(
old_state: "MoLeRDecoderState",
atom_symbol: str,
atom_logprob: float,
atom_choice_info: Optional[MoleculeGenerationAtomChoiceInfo] = None,
) -> "MoLeRDecoderState":
"""Add a new atom to the partial mocule under construction.
Note: This also sets the focus node to the next atom to consider, chosen from the
new ones added.
"""
if old_state._focus_atom is not None:
raise ValueError(
"New atoms can only be added when the decoder is not focused on creating bonds for an atom!"
)
new_mol = Chem.RWMol(old_state._molecule)
new_atom_idx = new_mol.AddAtom(initialise_atom_from_symbol(atom_symbol))
new_atom_types = list(old_state._atom_types)
new_atom_types.append(atom_symbol)
new_atoms_to_visit = list(old_state._atoms_to_visit)
new_atoms_to_visit.append(new_atom_idx)
new_focus_atom = new_atoms_to_visit.pop(0) # BFS exploration; .pop() would give DFS
new_generation_steps = MoLeRDecoderState.extend_generation_steps(
generation_steps=old_state._generation_steps,
choice_info=atom_choice_info,
molecule=new_mol,
)
return MoLeRDecoderState(
molecule_representation=old_state._molecule_representation,
molecule_id=old_state._molecule_id,
logprob=old_state._logprob + atom_logprob,
molecule=new_mol,
atom_types=new_atom_types,
adjacency_lists=old_state._adjacency_lists,
visited_atoms=old_state._visited_atoms,
atoms_to_visit=new_atoms_to_visit,
atoms_to_mark_as_visited=old_state._atoms_to_mark_as_visited,
focus_atom=new_focus_atom,
prior_focus_atom=old_state._focus_atom,
generation_steps=new_generation_steps,
candidate_attachment_points=old_state._candidate_attachment_points,
motifs=old_state._motifs,
num_free_bond_slots=old_state._num_free_bond_slots + [None],
)
@staticmethod
def new_for_finished_decoding(
old_state: "MoLeRDecoderState",
finish_logprob: float,
atom_choice_info: Optional[MoleculeGenerationAtomChoiceInfo] = None,
) -> "MoLeRDecoderState":
"""Create a new state for a finished decoding run."""
if old_state._focus_atom is not None:
raise ValueError(
"Decoding can only be finished when the decoder is not focused on creating bonds for an atom!"
)
new_generation_steps = MoLeRDecoderState.extend_generation_steps(
generation_steps=old_state._generation_steps,
choice_info=atom_choice_info,
molecule=old_state._molecule,
)
return MoLeRDecoderState(
molecule_representation=old_state._molecule_representation,
molecule_id=old_state._molecule_id,
logprob=old_state._logprob + finish_logprob,
molecule=old_state._molecule,
atom_types=old_state._atom_types,
adjacency_lists=old_state._adjacency_lists,
visited_atoms=old_state._visited_atoms,
atoms_to_visit=old_state._atoms_to_visit,
atoms_to_mark_as_visited=old_state._atoms_to_mark_as_visited,
focus_atom=-1,
prior_focus_atom=-1,
generation_steps=new_generation_steps,
candidate_attachment_points=old_state._candidate_attachment_points,
motifs=old_state._motifs,
num_free_bond_slots=old_state._num_free_bond_slots,
)
@staticmethod
def new_with_added_bond(
old_state: "MoLeRDecoderState",
target_atom_idx: int,
bond_type_idx: int,
bond_logprob: float,
edge_choice_info: Optional[MoleculeGenerationEdgeChoiceInfo] = None,
) -> "MoLeRDecoderState":
"""Add a new bond to the partial mocule under construction."""
if old_state._focus_atom is None:
raise ValueError(
"New bonds can only be added when the decoder is focused on creating bonds for an atom!"
)
new_mol = Chem.RWMol(old_state._molecule)
new_mol.AddBond(
old_state._focus_atom, target_atom_idx, EDGE_TYPE_IDX_TO_BOND_TYPE[bond_type_idx]
)
# We only need to make an actual copy of the adj_list for the bond type we change:
new_adjacency_lists = list(old_state._adjacency_lists)
new_adjacency_lists[bond_type_idx] = list(new_adjacency_lists[bond_type_idx])
new_adjacency_lists[bond_type_idx].append((old_state._focus_atom, target_atom_idx))
new_adjacency_lists[bond_type_idx].append((target_atom_idx, old_state._focus_atom))
new_generation_steps = MoLeRDecoderState.extend_generation_steps(
generation_steps=old_state._generation_steps,
choice_info=edge_choice_info,
molecule=new_mol,
)
if old_state._num_free_bond_slots[old_state._focus_atom] is not None:
raise ValueError("Focus atom has a constraint on the number of new bonds.")
if old_state._num_free_bond_slots[target_atom_idx] == 0:
raise ValueError("Tried to attach to an atom with no free bond slots.")
new_num_free_bond_slots = list(old_state._num_free_bond_slots)
# If we're connecting to an atom that has a constraint on the number of bonds, decrease that
# by 1. Note that we don't differentiate single/double/triple bonds here.
if new_num_free_bond_slots[target_atom_idx] is not None:
new_num_free_bond_slots[target_atom_idx] -= 1
return MoLeRDecoderState(
molecule_representation=old_state._molecule_representation,
molecule_id=old_state._molecule_id,
logprob=old_state._logprob + bond_logprob,
molecule=new_mol,
atom_types=old_state._atom_types,
adjacency_lists=new_adjacency_lists,
visited_atoms=old_state._visited_atoms,
atoms_to_visit=old_state._atoms_to_visit,
atoms_to_mark_as_visited=old_state._atoms_to_mark_as_visited,
focus_atom=old_state._focus_atom,
prior_focus_atom=old_state._prior_focus_atom,
generation_steps=new_generation_steps,
candidate_attachment_points=old_state._candidate_attachment_points,
motifs=old_state._motifs,
num_free_bond_slots=new_num_free_bond_slots,
)
@staticmethod
def new_with_added_motif(
old_state: "MoLeRDecoderState",
motif_type: str,
motif_logprob: float,
atom_choice_info: Optional[MoleculeGenerationAtomChoiceInfo] = None,
) -> "MoLeRDecoderState":
"""Add a new atom or entire motif to the partial mocule under construction."""
motif = Chem.MolFromSmiles(motif_type)
Chem.rdmolops.Kekulize(motif, clearAromaticFlags=True)
motif_atom_symbols = [get_atom_symbol(atom) for atom in motif.GetAtoms()]
motif_num_atoms = len(motif_atom_symbols)
node_idx_offset = len(old_state.molecule.GetAtoms())
motif_adjacency_list = [
(bond.GetBeginAtomIdx(), bond.GetEndAtomIdx(), BOND_DICT[str(bond.GetBondType())])
for bond in motif.GetBonds()
]
# Bonds to add to the motif, grouped by the larger of endpoint ids.
motif_bonds = [[] for _ in range(motif_num_atoms)]
for node_start, node_end, bond_type in motif_adjacency_list:
if node_start > node_end:
node_start, node_end = node_end, node_start
motif_bonds[node_end].append((node_start, bond_type))
assert not old_state._atoms_to_visit
assert old_state._focus_atom is None
current_state = old_state
for atom_symbol, bonds in zip(motif_atom_symbols, motif_bonds):
current_state._focus_atom = None
current_state = MoLeRDecoderState.new_with_added_atom(
old_state=current_state,
atom_symbol=atom_symbol,
atom_logprob=0.0, # We will add the full motif logprob at the end.
)
for target_atom_idx, bond_type in bonds:
current_state = MoLeRDecoderState.new_with_added_bond(
old_state=current_state,
target_atom_idx=node_idx_offset + target_atom_idx,
bond_type_idx=bond_type,
bond_logprob=0.0, # We will add the full motif logprob at the end
)
motif_nodes = list(range(node_idx_offset, node_idx_offset + motif_num_atoms))
motif_node_symmetry_class = Chem.CanonicalRankAtoms(motif, breakTies=False)
attachment_points = []
seen_symmetry_classes = set()
for node_idx, symmetry_class in enumerate(motif_node_symmetry_class):
if symmetry_class not in seen_symmetry_classes:
seen_symmetry_classes.add(symmetry_class)
attachment_points.append(node_idx)
attachment_points = get_open_attachment_points(
valid_attachment_points=attachment_points,
adjacency_list=[Edge(*args) for args in motif_adjacency_list],
node_types=motif_atom_symbols,
)
attachment_points = [node_idx + node_idx_offset for node_idx in attachment_points]
new_generation_steps = MoLeRDecoderState.extend_generation_steps(
generation_steps=old_state._generation_steps,
choice_info=atom_choice_info,
molecule=current_state._molecule,
)
# Record the freshly added motif, so it can be marked in the partial graph node features.
# Note that `symmetry_class_id` is not used here.
new_motifs = list(old_state._motifs)
new_motifs.append(
MotifAnnotation(
motif_type=motif_type,
atoms=[
MotifAtomAnnotation(atom_id=atom_id, symmetry_class_id=-1)
for atom_id in motif_nodes
],
)
)
return MoLeRDecoderState(
molecule_representation=current_state._molecule_representation,
molecule_id=old_state._molecule_id,
logprob=old_state._logprob + motif_logprob,
molecule=current_state._molecule,
atom_types=current_state._atom_types,
adjacency_lists=current_state._adjacency_lists,
visited_atoms=current_state._visited_atoms,
atoms_to_visit=None,
atoms_to_mark_as_visited=motif_nodes,
focus_atom=None,
prior_focus_atom=old_state._prior_focus_atom,
generation_steps=new_generation_steps,
candidate_attachment_points=attachment_points,
motifs=new_motifs,
num_free_bond_slots=old_state._num_free_bond_slots + [None] * motif_num_atoms,
)
@staticmethod
def new_with_focus_on_attachment_point(
old_state: "MoLeRDecoderState",
new_focus_atom: int,
focus_atom_logprob: float,
attachment_point_choice_info: MoleculeGenerationAttachmentPointChoiceInfo,
) -> "MoLeRDecoderState":
new_generation_steps = MoLeRDecoderState.extend_generation_steps(
generation_steps=old_state._generation_steps,
choice_info=attachment_point_choice_info,
molecule=old_state._molecule,
)
return MoLeRDecoderState(
molecule_representation=old_state._molecule_representation,
molecule_id=old_state._molecule_id,
logprob=old_state._logprob + focus_atom_logprob,
molecule=old_state._molecule,
atom_types=old_state._atom_types,
adjacency_lists=old_state._adjacency_lists,
visited_atoms=old_state._visited_atoms,
atoms_to_visit=old_state._atoms_to_visit,
atoms_to_mark_as_visited=old_state._atoms_to_mark_as_visited,
focus_atom=new_focus_atom,
prior_focus_atom=old_state._focus_atom,
generation_steps=new_generation_steps,
candidate_attachment_points=old_state._candidate_attachment_points,
motifs=old_state._motifs,
num_free_bond_slots=old_state._num_free_bond_slots,
)
@staticmethod
def new_with_focus_marked_as_visited(
old_state: "MoLeRDecoderState",
focus_node_finished_logprob: float,
edge_choice_info: Optional[MoleculeGenerationEdgeChoiceInfo] = None,
) -> "MoLeRDecoderState":
"""Mark a focus node as visited. If more nodes are available to visit, directly resets
the focus node to one of them.
"""
if old_state._focus_atom is None:
raise ValueError(
"Focus can only be marked as visited when the decoder is focused on creating bonds for an atom!"
)
new_visited_atoms = list(old_state._visited_atoms)
new_visited_atoms.append(old_state._focus_atom)
if len(old_state._atoms_to_visit) > 0:
new_atoms_to_visit = list(old_state._atoms_to_visit)
new_focus_atom: Optional[int] = new_atoms_to_visit.pop(
0
) # BFS exploration; .pop() would give DFS
else:
new_atoms_to_visit = old_state._atoms_to_visit
new_focus_atom = None
if old_state._atoms_to_mark_as_visited:
# The last node addition step added a motif - we now need to mark all motif nodes as visited.
count_already_in_visited = 0
for node_idx in old_state._atoms_to_mark_as_visited:
if node_idx in new_visited_atoms:
count_already_in_visited += 1
else:
new_visited_atoms.append(node_idx)
# The should be one motif node already added as visited: the attachment point.
assert count_already_in_visited == 1
new_generation_steps = MoLeRDecoderState.extend_generation_steps(
generation_steps=old_state._generation_steps,
choice_info=edge_choice_info,
molecule=old_state._molecule,
)
return MoLeRDecoderState(
molecule_representation=old_state._molecule_representation,
molecule_id=old_state._molecule_id,
logprob=old_state._logprob + focus_node_finished_logprob,
molecule=old_state._molecule,
atom_types=old_state._atom_types,
adjacency_lists=old_state._adjacency_lists,
visited_atoms=new_visited_atoms,
atoms_to_visit=new_atoms_to_visit,
atoms_to_mark_as_visited=None,
focus_atom=new_focus_atom,
prior_focus_atom=old_state._focus_atom,
generation_steps=new_generation_steps,
candidate_attachment_points=old_state._candidate_attachment_points,
motifs=old_state._motifs,
num_free_bond_slots=old_state._num_free_bond_slots,
)
@property
def molecule(self):
return self._molecule
@property
def molecule_representation(self):
return self._molecule_representation
@property
def molecule_id(self):
return self._molecule_id
@property
def logprob(self) -> float:
return self._logprob
@property
def focus_atom(self):
return self._focus_atom
@property
def prior_focus_atom(self):
return self._prior_focus_atom
@property
def adjacency_lists(self):
return self._adjacency_lists
@property
def atoms_to_visit(self):
return self._atoms_to_visit
@property
def atoms_to_mark_as_visited(self):
return self._atoms_to_mark_as_visited
@property
def generation_steps(self):
return self._generation_steps
@property
def candidate_attachment_points(self):
return self._candidate_attachment_points
def get_node_features(
self,
atom_featurisers: List[AtomFeatureExtractor],
motif_vocabulary: Optional[MotifVocabulary] = None,
):
"""
Compute node features for consumption in a GNN from the partial molecule
in this decoder state.
"""
# The overall computation here should not be cached, as we change the molecule
# between calls (adding atoms/bonds) and hence would never expect the same
# return values.
self._molecule.UpdatePropertyCache(strict=False)
features = featurise_atoms(
mol=self._molecule,
atom_feature_extractors=atom_featurisers,
motif_vocabulary=motif_vocabulary,
motifs=self._motifs,
)
return np.concatenate([features.real_valued_features]), features.categorical_features
def get_bond_candidate_targets(self) -> Tuple[np.ndarray, np.ndarray]:
"""
Compute candidates for bonds from the current focus atom.
Returns:
Pair of two numpy arrays.
First is a list of all valid targets of bonds from the current focus node. Shape: [CE]
Second is a mask of allowed bond types for the valid targets. Shape: [CE, 3]
"""
# We can only make bonds to nodes we have visited, but not exhausted the maximum number of
# bonds; this is our initial target set:
candidate_targets = {
idx for idx in self._visited_atoms if self._num_free_bond_slots[idx] != 0
}
# However, we remove those that we have already connected to:
for adj_list in self._adjacency_lists:
for edge in adj_list:
if edge[0] == self._focus_atom:
candidate_targets.discard(edge[1])
candidate_targets = np.array(list(candidate_targets), dtype=np.int32)
# Go from list of lists of tuples to list of numpy arrays, as we use slicing later:
adjacency_lists = [np.array(adj_list) for adj_list in self._adjacency_lists]
# Figure out which of these connections are chemically invalid:
candidate_targets_mask = constrain_edge_choices_based_on_valence(
start_node=self._focus_atom,
candidate_target_nodes=candidate_targets,
adjacency_lists=adjacency_lists,
node_types=self._atom_types,
)
bond_type_mask = constrain_edge_types_based_on_valence(
start_node=self._focus_atom,
candidate_target_nodes=candidate_targets,
# Go from list of lists of tuples to list of numpy arrays, as we use slicing inside:
adjacency_lists=adjacency_lists,
node_types=self._atom_types,
)
return candidate_targets[candidate_targets_mask], bond_type_mask[candidate_targets_mask]
def compute_bond_candidate_features(self, candidate_targets) -> np.ndarray:
"""
Compute features for bond candidates.
Note: This has to match
Args:
candidate_targets: NumPy array of bond targets from the current focus node.
Shape [EC], dtype int32
Returns:
NumPy array of features, where the i-th row of features corresponds to the
i-th bond candidate. Shape [EC, EFD], dtype int32
"""
# Compute distance features:
distance_features: List[int] = calculate_dist_from_focus_to_valid_target(
adjacency_list=self._adjacency_lists,
focus_node=self._focus_atom,
target_nodes=candidate_targets,
symmetrise_adjacency_list=False,
) # Shape [EC]
# Calculate the topology featues:
topology_features = calculate_topology_features(
edges=[(self._focus_atom, target) for target in candidate_targets], mol=self._molecule
) # Shape [EC, 2]
return np.concatenate(
[np.expand_dims(distance_features, axis=-1), topology_features], axis=1
)
def restrict_to_beam_size_per_mol(
decoder_states: Iterable[MoLeRDecoderState], beam_size: int
) -> List[MoLeRDecoderState]:
grouped_by_mol: DefaultDict[Any, List[MoLeRDecoderState]] = defaultdict(list)
for decoder_state in decoder_states:
grouped_by_mol[decoder_state.molecule_id].append(decoder_state)
restricted_decoder_states = []
for per_mol_decoder_states in grouped_by_mol.values():
restricted_decoder_states.extend(
heapq.nlargest(beam_size, per_mol_decoder_states, key=lambda s: s.logprob)
)
return restricted_decoder_states