extend SynthonCompletion for arbitrary node features

Author: KiddoZhu

torchdrug/data/graph.py

@@ -240,7 +240,7 @@ def connected_components(self):
             last = min_neighbor
             min_neighbor = scatter_min(min_neighbor[node_out], node_in, dim_size=self.num_node)[0]
         anchor = torch.unique(min_neighbor)
-        num_cc = scatter_add(torch.ones_like(anchor), self.node2graph[anchor])
+        num_cc = scatter_add(torch.ones_like(anchor), self.node2graph[anchor], dim_size=self.batch_size)
         return self.split(min_neighbor), num_cc
 
     def split(self, node2graph):

torchdrug/tasks/generation.py

@@ -99,8 +99,8 @@ def preprocess(self, train_set, valid_set, test_set):
             self.max_node = 0
 
             train_set = tqdm(train_set, "Computing max number of nodes and edge unrolling")
-            for data in train_set:
-                graph = data["graph"]
+            for sample in train_set:
+                graph = sample["graph"]
                 if graph.edge_list.numel():
                     edge_unroll = (graph.edge_list[:, 0] - graph.edge_list[:, 1]).abs().max().item()
                     self.max_edge_unroll = max(self.max_edge_unroll, edge_unroll)
@@ -677,8 +677,8 @@ def preprocess(self, train_set, valid_set, test_set):
             self.max_node = 0
 
             train_set = tqdm(train_set, "Computing max number of nodes and edge unrolling")
-            for data in train_set:
-                graph = data["graph"]
+            for sample in train_set:
+                graph = sample["graph"]
                 if graph.edge_list.numel():
                     edge_unroll = (graph.edge_list[:, 0] - graph.edge_list[:, 1]).abs().max().item()
                     self.max_edge_unroll = max(self.max_edge_unroll, edge_unroll)

torchdrug/tasks/retrosynthesis.py

@@ -1,3 +1,4 @@
+import inspect
 from collections import deque
 
 import torch
@@ -43,9 +44,9 @@ def __init__(self, model, feature=("reaction", "graph", "atom", "bond"), num_mlp
     def preprocess(self, train_set, valid_set, test_set):
         reaction_types = set()
         bond_types = set()
-        for data in train_set:
-            reaction_types.add(data["reaction"])
-            for graph in data["graph"]:
+        for sample in train_set:
+            reaction_types.add(sample["reaction"])
+            for graph in sample["graph"]:
                 bond_types.update(graph.edge_list[:, 2].tolist())
         self.num_reaction = len(reaction_types)
         self.num_relation = len(bond_types)
@@ -312,35 +313,22 @@ def __init__(self, model, feature=("reaction", "graph", "atom"), num_mlp_layer=2
 
     def preprocess(self, train_set, valid_set, test_set):
         reaction_types = set()
-        for data in train_set:
-            reaction_types.add(data["reaction"])
-        self.num_reaction = len(reaction_types)
-
-        if isinstance(train_set, torch_data.Subset):
-            dataset = train_set.dataset
-        else:
-            dataset = train_set
-        dataset.transform = transforms.Compose([
-            dataset.transform,
-            RandomBFSOrder(),
-        ])
-
-        # atom_types = set()
-        # bond_types = set()
-        # for data in train_set:
-        #     for graph in data["graph"]:
-        #         atom_types.update(graph.atom_type.tolist())
-        #         bond_types.update(graph.edge_list[:, 2].tolist())
-        # atom_types = torch.tensor(sorted(atom_types))
-
+        atom_types = set()
+        bond_types = set()
+        for sample in train_set:
+            reaction_types.add(sample["reaction"])
+            for graph in sample["graph"]:
+                atom_types.update(graph.atom_type.tolist())
+                bond_types.update(graph.edge_list[:, 2].tolist())
         # TODO: only for fast debugging, to remove
-        atom_types = torch.tensor([5, 6, 7, 8, 9, 12, 14, 15, 16, 17, 29, 30, 34, 35, 50, 53])
-        bond_types = torch.tensor([0, 1, 2])
-
+        # atom_types = torch.tensor([5, 6, 7, 8, 9, 12, 14, 15, 16, 17, 29, 30, 34, 35, 50, 53])
+        # bond_types = torch.tensor([0, 1, 2])
+        atom_types = torch.tensor(sorted(atom_types))
         atom2id = -torch.ones(atom_types.max() + 1, dtype=torch.long)
         atom2id[atom_types] = torch.arange(len(atom_types))
         self.register_buffer("id2atom", atom_types)
         self.register_buffer("atom2id", atom2id)
+        self.num_reaction = len(reaction_types)
         self.num_atom_type = len(atom_types)
         self.num_bond_type = len(bond_types)
         node_feature_dim = train_set[0]["graph"][0].node_feature.shape[-1]
@@ -349,7 +337,18 @@ def preprocess(self, train_set, valid_set, test_set):
             dataset = train_set.dataset
         else:
             dataset = train_set
-        self.dataset_kwargs = dataset.config_dict().get("kwargs")
+        dataset.transform = transforms.Compose([
+            dataset.transform,
+            RandomBFSOrder(),
+        ])
+        sig = inspect.signature(data.PackedMolecule.from_molecule)
+        keys = set(sig.parameters.keys())
+        kwargs = dataset.config_dict()
+        feature_kwargs = {}
+        for k, v in kwargs.items():
+            if k in keys:
+                feature_kwargs[k] = v
+        self.feature_kwargs = feature_kwargs
 
         node_dim = self.model.output_dim
         edge_dim = 0
@@ -382,7 +381,7 @@ def _update_molecule_feature(self, graphs):
         mols = graphs.to_molecule(ignore_error=True)
         valid = [mol is not None for mol in mols]
         valid = torch.tensor(valid, device=graphs.device)
-        new_graphs = type(graphs).from_molecule(mols, node_feature="synthon_completion", kekulize=True)
+        new_graphs = type(graphs).from_molecule(mols, **self.feature_kwargs)
 
         node_feature = torch.zeros(graphs.num_node, *new_graphs.node_feature.shape[1:],
                                    dtype=new_graphs.node_feature.dtype, device=graphs.device)
@@ -915,8 +914,7 @@ def predict_reactant(self, batch, num_beam=10, max_prediction=20, max_step=20):
             order = key.argsort(descending=True)
             new_graph = new_graph[order]
 
-            num_candidate = scatter_add(torch.ones_like(new_graph.synthon_id), new_graph.synthon_id,
-                                        dim_size=len(synthon))
+            num_candidate = new_graph.synthon_id.bincount(minlength=len(synthon))
             topk = functional.variadic_topk(new_graph.logp, num_candidate, num_beam)[1]
             topk_index = topk + (num_candidate.cumsum(0) - num_candidate).unsqueeze(-1)
             topk_index = torch.unique(topk_index)
@@ -965,7 +963,7 @@ def _extend(self, data, num_xs, input, input2graph=None):
             num_input_per_graph = len(input) // len(num_xs)
             input2graph = torch.arange(len(num_xs), device=data.device).unsqueeze(-1)
             input2graph = input2graph.repeat(1, num_input_per_graph).flatten()
-        num_inputs = scatter_add(torch.ones_like(input2graph), input2graph, dim_size=len(num_xs))
+        num_inputs = input2graph.bincount(minlength=len(num_xs))
         new_num_xs = num_xs + num_inputs
         new_num_cum_xs = new_num_xs.cumsum(0)
         new_num_x = new_num_cum_xs[-1].item()