torch-molecule 0.1.0.dev2__py3-none-any.whl
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.
- torch_molecule/__init__.py +70 -0
- torch_molecule/base/__init__.py +8 -0
- torch_molecule/base/base.py +396 -0
- torch_molecule/base/encoder.py +44 -0
- torch_molecule/base/generator.py +24 -0
- torch_molecule/base/old_encoder.py +322 -0
- torch_molecule/base/old_predictor.py +904 -0
- torch_molecule/base/old_predictor2.py +436 -0
- torch_molecule/base/predictor.py +112 -0
- torch_molecule/encoder/attrmask/__init__.py +3 -0
- torch_molecule/encoder/attrmask/model.py +113 -0
- torch_molecule/encoder/attrmask/modeling_attrmask.py +304 -0
- torch_molecule/encoder/constant.py +13 -0
- torch_molecule/encoder/contextpred/__init__.py +3 -0
- torch_molecule/encoder/contextpred/model.py +207 -0
- torch_molecule/encoder/contextpred/modeling_contextpred.py +299 -0
- torch_molecule/encoder/contextpred/utils.py +206 -0
- torch_molecule/encoder/edgepred/__init__.py +3 -0
- torch_molecule/encoder/edgepred/model.py +104 -0
- torch_molecule/encoder/edgepred/modeling_edgepred.py +285 -0
- torch_molecule/encoder/graphmae/__init__.py +3 -0
- torch_molecule/encoder/graphmae/dataloader.py +206 -0
- torch_molecule/encoder/graphmae/model.py +137 -0
- torch_molecule/encoder/graphmae/modeling_graphmae.py +323 -0
- torch_molecule/encoder/infograph/__init__.py +3 -0
- torch_molecule/encoder/infograph/loss_function.py +128 -0
- torch_molecule/encoder/infograph/model.py +92 -0
- torch_molecule/encoder/infograph/modeling_infograph.py +301 -0
- torch_molecule/encoder/infograph/module.py +31 -0
- torch_molecule/encoder/moama/__init__.py +3 -0
- torch_molecule/encoder/moama/model.py +131 -0
- torch_molecule/encoder/moama/modeling_moama.py +307 -0
- torch_molecule/encoder/moama/utils.py +143 -0
- torch_molecule/encoder/pretrained/__init__.py +3 -0
- torch_molecule/encoder/pretrained/modeling_pretrained.py +274 -0
- torch_molecule/encoder/supervised/__init__.py +3 -0
- torch_molecule/encoder/supervised/model.py +97 -0
- torch_molecule/encoder/supervised/modeling_supervised.py +354 -0
- torch_molecule/generator/digress/__init__.py +3 -0
- torch_molecule/generator/digress/diffusion.py +194 -0
- torch_molecule/generator/digress/modeling_digress.py +503 -0
- torch_molecule/generator/digress/transformer.py +360 -0
- torch_molecule/generator/digress/utils.py +196 -0
- torch_molecule/generator/gdss/__init__.py +3 -0
- torch_molecule/generator/gdss/layers.py +257 -0
- torch_molecule/generator/gdss/model.py +204 -0
- torch_molecule/generator/gdss/modeling_gdss.py +474 -0
- torch_molecule/generator/gdss/sde.py +334 -0
- torch_molecule/generator/gdss/solver.py +227 -0
- torch_molecule/generator/gdss/toremove/ScoreNetworkA.py +167 -0
- torch_molecule/generator/gdss/toremove/ScoreNetworkX.py +85 -0
- torch_molecule/generator/gdss/utils.py +240 -0
- torch_molecule/generator/graph_dit/__init__.py +3 -0
- torch_molecule/generator/graph_dit/diffusion.py +188 -0
- torch_molecule/generator/graph_dit/modeling_graph_dit.py +581 -0
- torch_molecule/generator/graph_dit/transformer.py +207 -0
- torch_molecule/generator/graph_dit/utils.py +220 -0
- torch_molecule/generator/graph_ga/__init__.py +3 -0
- torch_molecule/generator/graph_ga/crossover.py +182 -0
- torch_molecule/generator/graph_ga/modeling_graph_ga.py +310 -0
- torch_molecule/generator/graph_ga/mutate.py +136 -0
- torch_molecule/generator/graph_ga/oracle.py +95 -0
- torch_molecule/generator/jtvae/__init__.py +3 -0
- torch_molecule/generator/jtvae/jtnn/chemutils.py +450 -0
- torch_molecule/generator/jtvae/jtnn/datautils.py +127 -0
- torch_molecule/generator/jtvae/jtnn/jtmpn.py +140 -0
- torch_molecule/generator/jtvae/jtnn/jtnn_dec.py +350 -0
- torch_molecule/generator/jtvae/jtnn/jtnn_enc.py +131 -0
- torch_molecule/generator/jtvae/jtnn/mol_tree.py +111 -0
- torch_molecule/generator/jtvae/jtnn/mpn.py +129 -0
- torch_molecule/generator/jtvae/jtnn/nnutils.py +66 -0
- torch_molecule/generator/jtvae/jtnn/vocab.py +31 -0
- torch_molecule/generator/jtvae/jtnn_vae.py +246 -0
- torch_molecule/generator/jtvae/modeling_jtvae.py +247 -0
- torch_molecule/generator/lstm/__init__.py +3 -0
- torch_molecule/generator/lstm/action_sampler.py +91 -0
- torch_molecule/generator/lstm/lstm.py +59 -0
- torch_molecule/generator/lstm/modeling_lstm.py +263 -0
- torch_molecule/generator/lstm/smiles_char_dict.py +123 -0
- torch_molecule/generator/lstm/utils.py +62 -0
- torch_molecule/generator/molgpt/__init__.py +3 -0
- torch_molecule/generator/molgpt/dataset.py +87 -0
- torch_molecule/generator/molgpt/gpt.py +251 -0
- torch_molecule/generator/molgpt/modeling_molgpt.py +427 -0
- torch_molecule/generator/pretrained/modeling_pretrained.py +1 -0
- torch_molecule/nn/__init__.py +16 -0
- torch_molecule/nn/attention.py +63 -0
- torch_molecule/nn/embedder.py +183 -0
- torch_molecule/nn/gnn.py +631 -0
- torch_molecule/nn/mlp.py +34 -0
- torch_molecule/predictor/bfgnn/__init__.py +3 -0
- torch_molecule/predictor/bfgnn/model.py +116 -0
- torch_molecule/predictor/bfgnn/modeling_bfgnn.py +81 -0
- torch_molecule/predictor/dir/__init__.py +3 -0
- torch_molecule/predictor/dir/model.py +235 -0
- torch_molecule/predictor/dir/modeling_dir.py +161 -0
- torch_molecule/predictor/dir/utils.py +40 -0
- torch_molecule/predictor/gnn/__init__.py +3 -0
- torch_molecule/predictor/gnn/model.py +114 -0
- torch_molecule/predictor/gnn/modeling_gnn.py +706 -0
- torch_molecule/predictor/gnn/old_modeling_gnn.py +707 -0
- torch_molecule/predictor/grea/__init__.py +3 -0
- torch_molecule/predictor/grea/model.py +190 -0
- torch_molecule/predictor/grea/modeling_grea.py +107 -0
- torch_molecule/predictor/grea/old1_features.py +229 -0
- torch_molecule/predictor/grea/old1_modeling_grea.py +493 -0
- torch_molecule/predictor/grea/utils.py +97 -0
- torch_molecule/predictor/grin/__init__.py +3 -0
- torch_molecule/predictor/grin/model.py +116 -0
- torch_molecule/predictor/grin/modeling_grin.py +91 -0
- torch_molecule/predictor/irm/__init__.py +3 -0
- torch_molecule/predictor/irm/model.py +127 -0
- torch_molecule/predictor/irm/modeling_irm.py +179 -0
- torch_molecule/predictor/lstm/__init__.py +3 -0
- torch_molecule/predictor/lstm/model.py +145 -0
- torch_molecule/predictor/lstm/modeling_lstm.py +627 -0
- torch_molecule/predictor/lstm/token_from_smiles.py +85 -0
- torch_molecule/predictor/old_modeling_graph.py +650 -0
- torch_molecule/predictor/rpgnn/__init__.py +3 -0
- torch_molecule/predictor/rpgnn/model.py +185 -0
- torch_molecule/predictor/rpgnn/modeling_rpgnn.py +76 -0
- torch_molecule/predictor/sgir/__init__.py +3 -0
- torch_molecule/predictor/sgir/modeling_sgir.py +260 -0
- torch_molecule/predictor/sgir/strategy.py +266 -0
- torch_molecule/predictor/smiles_transformer/__init__.py +3 -0
- torch_molecule/predictor/smiles_transformer/model.py +424 -0
- torch_molecule/predictor/smiles_transformer/modeling_transformer.py +252 -0
- torch_molecule/predictor/ssr/__init__.py +3 -0
- torch_molecule/predictor/ssr/model.py +283 -0
- torch_molecule/predictor/ssr/modeling_ssr.py +267 -0
- torch_molecule/predictor/ssr/utils.py +185 -0
- torch_molecule/utils/__init__.py +31 -0
- torch_molecule/utils/checker.py +158 -0
- torch_molecule/utils/checkpoint.py +294 -0
- torch_molecule/utils/format.py +140 -0
- torch_molecule/utils/generic/metrics.py +530 -0
- torch_molecule/utils/generic/pseudo_tasks.py +11 -0
- torch_molecule/utils/generic/weights.py +30 -0
- torch_molecule/utils/graph/features.py +178 -0
- torch_molecule/utils/graph/graph_from_smiles.py +134 -0
- torch_molecule/utils/graph/graph_to_smiles.py +312 -0
- torch_molecule/utils/hf.py +232 -0
- torch_molecule/utils/search.py +64 -0
- torch_molecule-0.1.0.dev2.dist-info/METADATA +243 -0
- torch_molecule-0.1.0.dev2.dist-info/RECORD +148 -0
- torch_molecule-0.1.0.dev2.dist-info/WHEEL +5 -0
- torch_molecule-0.1.0.dev2.dist-info/licenses/LICENSE +21 -0
- torch_molecule-0.1.0.dev2.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
|
|
4
|
+
|
|
5
|
+
from ...nn import GNN_node, GNN_node_Virtualnode, MLP
|
|
6
|
+
from ...utils import init_weights
|
|
7
|
+
|
|
8
|
+
class GRIN(nn.Module):
|
|
9
|
+
def __init__(
|
|
10
|
+
self,
|
|
11
|
+
num_task,
|
|
12
|
+
num_layer,
|
|
13
|
+
hidden_size=300,
|
|
14
|
+
gnn_type="gin-virtual",
|
|
15
|
+
drop_ratio=0.5,
|
|
16
|
+
norm_layer="batch_norm",
|
|
17
|
+
graph_pooling="max",
|
|
18
|
+
augmented_feature=['maccs', 'morgan'],
|
|
19
|
+
algorithm_aligned = 'mst'
|
|
20
|
+
):
|
|
21
|
+
super(GRIN, self).__init__()
|
|
22
|
+
gnn_name = gnn_type.split("-")[0]
|
|
23
|
+
self.algorithm_aligned = algorithm_aligned
|
|
24
|
+
self.num_task = num_task
|
|
25
|
+
self.hidden_size = hidden_size
|
|
26
|
+
|
|
27
|
+
if "virtual" in gnn_type:
|
|
28
|
+
self.graph_encoder = GNN_node_Virtualnode(
|
|
29
|
+
num_layer,
|
|
30
|
+
hidden_size,
|
|
31
|
+
JK="last",
|
|
32
|
+
drop_ratio=drop_ratio,
|
|
33
|
+
residual=True,
|
|
34
|
+
gnn_name=gnn_name,
|
|
35
|
+
norm_layer=norm_layer,
|
|
36
|
+
algorithm_aligned=algorithm_aligned
|
|
37
|
+
)
|
|
38
|
+
else:
|
|
39
|
+
self.graph_encoder = GNN_node(
|
|
40
|
+
num_layer,
|
|
41
|
+
hidden_size,
|
|
42
|
+
JK="last",
|
|
43
|
+
drop_ratio=drop_ratio,
|
|
44
|
+
residual=True,
|
|
45
|
+
gnn_name=gnn_name,
|
|
46
|
+
norm_layer=norm_layer,
|
|
47
|
+
algorithm_aligned=algorithm_aligned
|
|
48
|
+
)
|
|
49
|
+
if graph_pooling == "sum":
|
|
50
|
+
self.pool = global_add_pool
|
|
51
|
+
elif graph_pooling == "mean":
|
|
52
|
+
self.pool = global_mean_pool
|
|
53
|
+
elif graph_pooling == "max":
|
|
54
|
+
self.pool = global_max_pool
|
|
55
|
+
else:
|
|
56
|
+
raise ValueError(f"Invalid graph pooling type {graph_pooling}.")
|
|
57
|
+
|
|
58
|
+
graph_dim = hidden_size
|
|
59
|
+
self.augmented_feature = augmented_feature
|
|
60
|
+
if augmented_feature:
|
|
61
|
+
if "morgan" in augmented_feature:
|
|
62
|
+
graph_dim += 1024
|
|
63
|
+
if "maccs" in augmented_feature:
|
|
64
|
+
graph_dim += 167
|
|
65
|
+
self.predictor = MLP(graph_dim, hidden_features=2 * hidden_size, out_features=num_task)
|
|
66
|
+
|
|
67
|
+
def initialize_parameters(self, seed=None):
|
|
68
|
+
"""
|
|
69
|
+
Randomly initialize all model parameters using the init_weights function.
|
|
70
|
+
|
|
71
|
+
Args:
|
|
72
|
+
seed (int, optional): Random seed for reproducibility. Defaults to None.
|
|
73
|
+
"""
|
|
74
|
+
if seed is not None:
|
|
75
|
+
torch.manual_seed(seed)
|
|
76
|
+
|
|
77
|
+
# Initialize the main components
|
|
78
|
+
init_weights(self.graph_encoder)
|
|
79
|
+
init_weights(self.predictor)
|
|
80
|
+
|
|
81
|
+
# Reset all parameters using PyTorch Geometric's reset function
|
|
82
|
+
def reset_parameters(module):
|
|
83
|
+
if hasattr(module, 'reset_parameters'):
|
|
84
|
+
module.reset_parameters()
|
|
85
|
+
elif hasattr(module, 'weight') and hasattr(module.weight, 'data'):
|
|
86
|
+
init_weights(module)
|
|
87
|
+
|
|
88
|
+
self.apply(reset_parameters)
|
|
89
|
+
|
|
90
|
+
def _augmented_graph_features(self, batched_data, h_rep):
|
|
91
|
+
if self.augmented_feature:
|
|
92
|
+
if 'morgan' in self.augmented_feature:
|
|
93
|
+
morgan = batched_data.morgan.type_as(h_rep)
|
|
94
|
+
h_rep = torch.cat((h_rep, morgan), dim=1)
|
|
95
|
+
if 'maccs' in self.augmented_feature:
|
|
96
|
+
maccs = batched_data.maccs.type_as(h_rep)
|
|
97
|
+
h_rep = torch.cat((h_rep, maccs), dim=1)
|
|
98
|
+
return h_rep
|
|
99
|
+
|
|
100
|
+
def compute_loss(self, batched_data, criterion, l1_penalty=1e-3):
|
|
101
|
+
h_node, _ = self.graph_encoder(batched_data)
|
|
102
|
+
h_rep = self.pool(h_node, batched_data.batch)
|
|
103
|
+
h_rep = self._augmented_graph_features(batched_data, h_rep)
|
|
104
|
+
prediction = self.predictor(h_rep)
|
|
105
|
+
target = batched_data.y.to(torch.float32)
|
|
106
|
+
is_labeled = batched_data.y == batched_data.y
|
|
107
|
+
loss = criterion(prediction.to(torch.float32)[is_labeled], target[is_labeled]).mean()
|
|
108
|
+
l1_loss = sum(p.abs().sum() for p in self.graph_encoder.parameters()) + sum(p.abs().sum() for p in self.predictor.parameters())
|
|
109
|
+
return loss + l1_penalty * l1_loss
|
|
110
|
+
|
|
111
|
+
def forward(self, batched_data):
|
|
112
|
+
h_node, _ = self.graph_encoder(batched_data)
|
|
113
|
+
h_rep = self.pool(h_node, batched_data.batch)
|
|
114
|
+
h_rep = self._augmented_graph_features(batched_data, h_rep)
|
|
115
|
+
prediction = self.predictor(h_rep)
|
|
116
|
+
return {"prediction": prediction}
|
|
@@ -0,0 +1,91 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import numpy as np
|
|
3
|
+
import warnings
|
|
4
|
+
import datetime
|
|
5
|
+
from tqdm import tqdm
|
|
6
|
+
from typing import Optional, Union, Dict, Any, List, Type
|
|
7
|
+
from dataclasses import dataclass, field
|
|
8
|
+
|
|
9
|
+
import torch
|
|
10
|
+
from torch_geometric.data import Data
|
|
11
|
+
|
|
12
|
+
from .model import GRIN
|
|
13
|
+
from ...utils import graph_from_smiles
|
|
14
|
+
from ..gnn.modeling_gnn import GNNMolecularPredictor
|
|
15
|
+
from ...utils.search import (
|
|
16
|
+
ParameterSpec,
|
|
17
|
+
ParameterType,
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
@dataclass
|
|
21
|
+
class GRINMolecularPredictor(GNNMolecularPredictor):
|
|
22
|
+
"""
|
|
23
|
+
This predictor implements GRIN for Max Spanning Tree algorithm aligned GNN.
|
|
24
|
+
|
|
25
|
+
The full name of GRIN is Graph Invariant Representation Learning.
|
|
26
|
+
|
|
27
|
+
References
|
|
28
|
+
----------
|
|
29
|
+
- Learning Repetition-Invariant Representations for Polymer Informatics.
|
|
30
|
+
https://arxiv.org/pdf/2505.10726
|
|
31
|
+
|
|
32
|
+
:param l1_penalty: Weight for the L1 penalty
|
|
33
|
+
:type l1_penalty: float, default=1e-3
|
|
34
|
+
:param epochs_to_penalize: Number of epochs to train before starting L1 penalty
|
|
35
|
+
:type epochs_to_penalize: int, default=100
|
|
36
|
+
"""
|
|
37
|
+
|
|
38
|
+
l1_penalty: float = 1e-3
|
|
39
|
+
epochs_to_penalize: int = 100
|
|
40
|
+
|
|
41
|
+
# Other Non-init fields
|
|
42
|
+
model_name: str = "GRINMolecularPredictor"
|
|
43
|
+
model_class: Type[GRIN] = field(default=GRIN, init=False)
|
|
44
|
+
|
|
45
|
+
def __post_init__(self):
|
|
46
|
+
super().__post_init__()
|
|
47
|
+
|
|
48
|
+
@staticmethod
|
|
49
|
+
def _get_param_names() -> List[str]:
|
|
50
|
+
return GNNMolecularPredictor._get_param_names() + [
|
|
51
|
+
"l1_penalty",
|
|
52
|
+
"epochs_to_penalize"
|
|
53
|
+
]
|
|
54
|
+
def _get_default_search_space(self):
|
|
55
|
+
search_space = super()._get_default_search_space().copy()
|
|
56
|
+
search_space["l1_penalty"] = ParameterSpec(ParameterType.LOG_FLOAT, (1e-6, 1))
|
|
57
|
+
search_space["epochs_to_penalize"] = ParameterSpec(ParameterType.INTEGER, (0, 100))
|
|
58
|
+
return search_space
|
|
59
|
+
|
|
60
|
+
def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
|
|
61
|
+
base_params = super()._get_model_params(checkpoint)
|
|
62
|
+
return base_params
|
|
63
|
+
|
|
64
|
+
def _train_epoch(self, train_loader, optimizer, epoch):
|
|
65
|
+
self.model.train()
|
|
66
|
+
losses = []
|
|
67
|
+
|
|
68
|
+
iterator = (
|
|
69
|
+
tqdm(train_loader, desc="Training", leave=False)
|
|
70
|
+
if self.verbose
|
|
71
|
+
else train_loader
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
for step, batch in enumerate(iterator):
|
|
75
|
+
batch = batch.to(self.device)
|
|
76
|
+
optimizer.zero_grad()
|
|
77
|
+
if epoch >= self.epochs_to_penalize:
|
|
78
|
+
l1_penalty = min(epoch - self.epochs_to_penalize, 1) * self.l1_penalty
|
|
79
|
+
else:
|
|
80
|
+
l1_penalty = 0
|
|
81
|
+
loss = self.model.compute_loss(batch, self.loss_criterion, l1_penalty)
|
|
82
|
+
loss.backward()
|
|
83
|
+
if self.grad_clip_value is not None:
|
|
84
|
+
torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_value)
|
|
85
|
+
optimizer.step()
|
|
86
|
+
|
|
87
|
+
if self.verbose:
|
|
88
|
+
iterator.set_postfix({"Epoch": epoch, "Total Loss": f"{loss.item():.4f}"})
|
|
89
|
+
losses.append(loss.item())
|
|
90
|
+
|
|
91
|
+
return losses
|
|
@@ -0,0 +1,127 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
|
|
4
|
+
|
|
5
|
+
from ...nn import GNN_node, GNN_node_Virtualnode, MLP
|
|
6
|
+
from ...utils import init_weights
|
|
7
|
+
|
|
8
|
+
class GNN(nn.Module):
|
|
9
|
+
def __init__(
|
|
10
|
+
self,
|
|
11
|
+
num_task,
|
|
12
|
+
num_layer,
|
|
13
|
+
hidden_size=300,
|
|
14
|
+
gnn_type="gin-virtual",
|
|
15
|
+
drop_ratio=0.5,
|
|
16
|
+
norm_layer="batch_norm",
|
|
17
|
+
graph_pooling="max",
|
|
18
|
+
augmented_feature=['maccs', 'morgan']
|
|
19
|
+
):
|
|
20
|
+
super(GNN, self).__init__()
|
|
21
|
+
gnn_name = gnn_type.split("-")[0]
|
|
22
|
+
self.num_task = num_task
|
|
23
|
+
self.hidden_size = hidden_size
|
|
24
|
+
|
|
25
|
+
if "virtual" in gnn_type:
|
|
26
|
+
self.graph_encoder = GNN_node_Virtualnode(
|
|
27
|
+
num_layer,
|
|
28
|
+
hidden_size,
|
|
29
|
+
JK="last",
|
|
30
|
+
drop_ratio=drop_ratio,
|
|
31
|
+
residual=True,
|
|
32
|
+
gnn_name=gnn_name,
|
|
33
|
+
norm_layer=norm_layer,
|
|
34
|
+
)
|
|
35
|
+
else:
|
|
36
|
+
self.graph_encoder = GNN_node(
|
|
37
|
+
num_layer,
|
|
38
|
+
hidden_size,
|
|
39
|
+
JK="last",
|
|
40
|
+
drop_ratio=drop_ratio,
|
|
41
|
+
residual=True,
|
|
42
|
+
gnn_name=gnn_name,
|
|
43
|
+
norm_layer=norm_layer,
|
|
44
|
+
)
|
|
45
|
+
if graph_pooling == "sum":
|
|
46
|
+
self.pool = global_add_pool
|
|
47
|
+
elif graph_pooling == "mean":
|
|
48
|
+
self.pool = global_mean_pool
|
|
49
|
+
elif graph_pooling == "max":
|
|
50
|
+
self.pool = global_max_pool
|
|
51
|
+
else:
|
|
52
|
+
raise ValueError(f"Invalid graph pooling type {graph_pooling}.")
|
|
53
|
+
|
|
54
|
+
graph_dim = hidden_size
|
|
55
|
+
self.augmented_feature = augmented_feature
|
|
56
|
+
if augmented_feature:
|
|
57
|
+
if "morgan" in augmented_feature:
|
|
58
|
+
graph_dim += 1024
|
|
59
|
+
if "maccs" in augmented_feature:
|
|
60
|
+
graph_dim += 167
|
|
61
|
+
self.predictor = MLP(graph_dim, hidden_features=2 * hidden_size, out_features=num_task)
|
|
62
|
+
|
|
63
|
+
def initialize_parameters(self, seed=None):
|
|
64
|
+
"""
|
|
65
|
+
Randomly initialize all model parameters using the init_weights function.
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
seed (int, optional): Random seed for reproducibility. Defaults to None.
|
|
69
|
+
"""
|
|
70
|
+
if seed is not None:
|
|
71
|
+
torch.manual_seed(seed)
|
|
72
|
+
|
|
73
|
+
# Initialize the main components
|
|
74
|
+
init_weights(self.graph_encoder)
|
|
75
|
+
init_weights(self.predictor)
|
|
76
|
+
|
|
77
|
+
# Reset all parameters using PyTorch Geometric's reset function
|
|
78
|
+
def reset_parameters(module):
|
|
79
|
+
if hasattr(module, 'reset_parameters'):
|
|
80
|
+
module.reset_parameters()
|
|
81
|
+
elif hasattr(module, 'weight') and hasattr(module.weight, 'data'):
|
|
82
|
+
init_weights(module)
|
|
83
|
+
|
|
84
|
+
self.apply(reset_parameters)
|
|
85
|
+
|
|
86
|
+
def _augmented_graph_features(self, batched_data, h_rep):
|
|
87
|
+
if self.augmented_feature:
|
|
88
|
+
if 'morgan' in self.augmented_feature:
|
|
89
|
+
morgan = batched_data.morgan.type_as(h_rep)
|
|
90
|
+
h_rep = torch.cat((h_rep, morgan), dim=1)
|
|
91
|
+
if 'maccs' in self.augmented_feature:
|
|
92
|
+
maccs = batched_data.maccs.type_as(h_rep)
|
|
93
|
+
h_rep = torch.cat((h_rep, maccs), dim=1)
|
|
94
|
+
return h_rep
|
|
95
|
+
|
|
96
|
+
def compute_loss(self, batched_data, criterion, scale=1.0, penalty_weight=1.0):
|
|
97
|
+
h_node, _ = self.graph_encoder(batched_data)
|
|
98
|
+
h_rep = self.pool(h_node, batched_data.batch)
|
|
99
|
+
h_rep = self._augmented_graph_features(batched_data, h_rep)
|
|
100
|
+
prediction = self.predictor(h_rep)
|
|
101
|
+
target = batched_data.y.to(torch.float32)
|
|
102
|
+
is_labeled = batched_data.y == batched_data.y
|
|
103
|
+
|
|
104
|
+
dummy = torch.nn.Parameter(torch.Tensor([scale])).to(prediction.device)
|
|
105
|
+
losses_erm = criterion(prediction.to(torch.float32)[is_labeled] * dummy, target[is_labeled])
|
|
106
|
+
|
|
107
|
+
environments = batched_data.environment[is_labeled]
|
|
108
|
+
unique_envs = environments.unique()
|
|
109
|
+
env_losses = []
|
|
110
|
+
for env in unique_envs:
|
|
111
|
+
env_mask = environments == env
|
|
112
|
+
env_loss = losses_erm[env_mask].mean()
|
|
113
|
+
env_grad = torch.autograd.grad(env_loss, dummy, create_graph=True)[0]
|
|
114
|
+
env_losses.append(torch.sum(env_grad**2))
|
|
115
|
+
|
|
116
|
+
penalty = torch.sum(torch.stack(env_losses))
|
|
117
|
+
total_loss = losses_erm.mean() + penalty_weight * penalty
|
|
118
|
+
if penalty_weight > 1.0:
|
|
119
|
+
total_loss /= penalty_weight
|
|
120
|
+
return total_loss, losses_erm.mean(), penalty
|
|
121
|
+
|
|
122
|
+
def forward(self, batched_data):
|
|
123
|
+
h_node, _ = self.graph_encoder(batched_data)
|
|
124
|
+
h_rep = self.pool(h_node, batched_data.batch)
|
|
125
|
+
h_rep = self._augmented_graph_features(batched_data, h_rep)
|
|
126
|
+
prediction = self.predictor(h_rep)
|
|
127
|
+
return {"prediction": prediction}
|
|
@@ -0,0 +1,179 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import numpy as np
|
|
3
|
+
import warnings
|
|
4
|
+
import datetime
|
|
5
|
+
from tqdm import tqdm
|
|
6
|
+
from typing import Optional, Union, Dict, Any, List, Type
|
|
7
|
+
from dataclasses import dataclass, field
|
|
8
|
+
|
|
9
|
+
import torch
|
|
10
|
+
from torch_geometric.data import Data
|
|
11
|
+
|
|
12
|
+
from .model import GNN
|
|
13
|
+
from ...utils import graph_from_smiles
|
|
14
|
+
from ..gnn.modeling_gnn import GNNMolecularPredictor
|
|
15
|
+
from ...utils.search import (
|
|
16
|
+
ParameterSpec,
|
|
17
|
+
ParameterType,
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
@dataclass
|
|
21
|
+
class IRMMolecularPredictor(GNNMolecularPredictor):
|
|
22
|
+
"""This predictor implements a Invariant Risk Minimization model with the GNN.
|
|
23
|
+
|
|
24
|
+
The full name of IRM is Invariant Risk Minimization.
|
|
25
|
+
|
|
26
|
+
References
|
|
27
|
+
----------
|
|
28
|
+
- Invariant Risk Minimization.
|
|
29
|
+
https://arxiv.org/abs/1907.02893
|
|
30
|
+
|
|
31
|
+
- Reference Code: https://github.com/facebookresearch/InvariantRiskMinimization
|
|
32
|
+
|
|
33
|
+
Parameters
|
|
34
|
+
----------
|
|
35
|
+
IRM_environment : Union[torch.Tensor, np.ndarray, List, str], default="random"
|
|
36
|
+
Environment assignments for IRM. Can be a list of integers (one per sample),
|
|
37
|
+
or "random" to assign environments randomly.
|
|
38
|
+
scale : float, default=1.0
|
|
39
|
+
Scaling factor for the IRM penalty term.
|
|
40
|
+
penalty_weight : float, default=1.0
|
|
41
|
+
Weight of the IRM penalty in the loss function.
|
|
42
|
+
penalty_anneal_iters : int, default=100
|
|
43
|
+
Number of iterations for annealing the penalty weight.
|
|
44
|
+
"""
|
|
45
|
+
|
|
46
|
+
IRM_environment: Union[torch.Tensor, np.ndarray, List, str] = "random"
|
|
47
|
+
scale: float = 1.0
|
|
48
|
+
penalty_weight: float = 1.0
|
|
49
|
+
penalty_anneal_iters: int = 100
|
|
50
|
+
|
|
51
|
+
# Other Non-init fields
|
|
52
|
+
model_name: str = "IRMMolecularPredictor"
|
|
53
|
+
model_class: Type[GNN] = field(default=GNN, init=False)
|
|
54
|
+
|
|
55
|
+
def __post_init__(self):
|
|
56
|
+
super().__post_init__()
|
|
57
|
+
|
|
58
|
+
@staticmethod
|
|
59
|
+
def _get_param_names() -> List[str]:
|
|
60
|
+
return GNNMolecularPredictor._get_param_names() + [
|
|
61
|
+
"IRM_environment",
|
|
62
|
+
"scale",
|
|
63
|
+
"penalty_weight",
|
|
64
|
+
"penalty_anneal_iters",
|
|
65
|
+
]
|
|
66
|
+
|
|
67
|
+
def _get_default_search_space(self):
|
|
68
|
+
search_space = super()._get_default_search_space().copy()
|
|
69
|
+
search_space["penalty_weight"] = ParameterSpec(ParameterType.LOG_FLOAT, (1e-10, 1))
|
|
70
|
+
search_space["penalty_anneal_iters"] = ParameterSpec(ParameterType.INTEGER, (10, 100))
|
|
71
|
+
return search_space
|
|
72
|
+
|
|
73
|
+
def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
|
|
74
|
+
base_params = super()._get_model_params(checkpoint)
|
|
75
|
+
return base_params
|
|
76
|
+
|
|
77
|
+
def set_IRM_environment(self, environment: Union[torch.Tensor, np.ndarray, List, str]):
|
|
78
|
+
if isinstance(environment, str):
|
|
79
|
+
if environment != "random":
|
|
80
|
+
raise ValueError("IRM_environment must be 'random' if specified with a string")
|
|
81
|
+
self.IRM_environment = environment
|
|
82
|
+
else:
|
|
83
|
+
if isinstance(environment, np.ndarray) or isinstance(environment, torch.Tensor):
|
|
84
|
+
self.IRM_environment = environment.reshape(-1).tolist()
|
|
85
|
+
else:
|
|
86
|
+
self.IRM_environment = environment
|
|
87
|
+
|
|
88
|
+
if not all(isinstance(item, int) for item in self.IRM_environment):
|
|
89
|
+
raise ValueError("IRM_environment must be a list of integers")
|
|
90
|
+
|
|
91
|
+
def _convert_to_pytorch_data(self, X, y=None):
|
|
92
|
+
"""Convert numpy arrays to PyTorch Geometric data format.
|
|
93
|
+
"""
|
|
94
|
+
if self.verbose:
|
|
95
|
+
iterator = tqdm(enumerate(X), desc="Converting molecules to graphs", total=len(X))
|
|
96
|
+
else:
|
|
97
|
+
iterator = enumerate(X)
|
|
98
|
+
|
|
99
|
+
pyg_graph_list = []
|
|
100
|
+
for idx, smiles_or_mol in iterator:
|
|
101
|
+
if y is not None:
|
|
102
|
+
properties = y[idx]
|
|
103
|
+
else:
|
|
104
|
+
properties = None
|
|
105
|
+
graph = graph_from_smiles(smiles_or_mol, properties, self.augmented_feature)
|
|
106
|
+
g = Data()
|
|
107
|
+
g.num_nodes = graph["num_nodes"]
|
|
108
|
+
g.edge_index = torch.from_numpy(graph["edge_index"])
|
|
109
|
+
|
|
110
|
+
del graph["num_nodes"]
|
|
111
|
+
del graph["edge_index"]
|
|
112
|
+
|
|
113
|
+
if graph["edge_feat"] is not None:
|
|
114
|
+
g.edge_attr = torch.from_numpy(graph["edge_feat"])
|
|
115
|
+
del graph["edge_feat"]
|
|
116
|
+
|
|
117
|
+
if graph["node_feat"] is not None:
|
|
118
|
+
g.x = torch.from_numpy(graph["node_feat"])
|
|
119
|
+
del graph["node_feat"]
|
|
120
|
+
|
|
121
|
+
if graph["y"] is not None:
|
|
122
|
+
g.y = torch.from_numpy(graph["y"])
|
|
123
|
+
del graph["y"]
|
|
124
|
+
|
|
125
|
+
if graph["morgan"] is not None:
|
|
126
|
+
g.morgan = torch.tensor(graph["morgan"], dtype=torch.int8).view(1, -1)
|
|
127
|
+
del graph["morgan"]
|
|
128
|
+
|
|
129
|
+
if graph["maccs"] is not None:
|
|
130
|
+
g.maccs = torch.tensor(graph["maccs"], dtype=torch.int8).view(1, -1)
|
|
131
|
+
del graph["maccs"]
|
|
132
|
+
|
|
133
|
+
if self.IRM_environment == "random":
|
|
134
|
+
g.environment = torch.randint(0, 2, (1,)).view(1, 1)
|
|
135
|
+
elif len(X) != len(self.IRM_environment):
|
|
136
|
+
raise ValueError("IRM_environment must has the same length as the input, which is {}".format(len(X)))
|
|
137
|
+
else:
|
|
138
|
+
if isinstance(self.IRM_environment[idx], int):
|
|
139
|
+
g.environment = torch.tensor(self.IRM_environment[idx], dtype=torch.int64).view(1, 1)
|
|
140
|
+
else:
|
|
141
|
+
raise ValueError("IRM_environment must be a list of integers")
|
|
142
|
+
pyg_graph_list.append(g)
|
|
143
|
+
|
|
144
|
+
return pyg_graph_list
|
|
145
|
+
|
|
146
|
+
def _train_epoch(self, train_loader, optimizer, epoch):
|
|
147
|
+
self.model.train()
|
|
148
|
+
losses = []
|
|
149
|
+
losses_erm = []
|
|
150
|
+
penalties = []
|
|
151
|
+
|
|
152
|
+
iterator = (
|
|
153
|
+
tqdm(train_loader, desc="Training", leave=False)
|
|
154
|
+
if self.verbose
|
|
155
|
+
else train_loader
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
for step, batch in enumerate(iterator):
|
|
159
|
+
batch = batch.to(self.device)
|
|
160
|
+
optimizer.zero_grad()
|
|
161
|
+
|
|
162
|
+
if epoch >= self.penalty_anneal_iters:
|
|
163
|
+
penalty_weight = self.penalty_weight
|
|
164
|
+
else:
|
|
165
|
+
penalty_weight = 1.0
|
|
166
|
+
loss, loss_erm, penalty = self.model.compute_loss(batch, self.loss_criterion, self.scale, penalty_weight)
|
|
167
|
+
loss.backward()
|
|
168
|
+
if self.grad_clip_value is not None:
|
|
169
|
+
torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_value)
|
|
170
|
+
|
|
171
|
+
optimizer.step()
|
|
172
|
+
losses.append(loss.item())
|
|
173
|
+
losses_erm.append(loss_erm.item())
|
|
174
|
+
penalties.append(penalty.item())
|
|
175
|
+
|
|
176
|
+
if self.verbose:
|
|
177
|
+
iterator.set_postfix({"Epoch": epoch, "Total Loss": f"{loss.item():.4f}", "ERM Loss": f"{loss_erm.item():.4f}", "IRM Penalty": f"{penalty.item():.4f}"})
|
|
178
|
+
|
|
179
|
+
return losses
|
|
@@ -0,0 +1,145 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
from sklearn.metrics import mean_absolute_error, r2_score, mean_squared_error
|
|
4
|
+
|
|
5
|
+
from ...utils import init_weights
|
|
6
|
+
|
|
7
|
+
# Define the PyTorch-based LSTM model
|
|
8
|
+
class LSTM(nn.Module):
|
|
9
|
+
"""LSTM-based model for molecular property prediction.
|
|
10
|
+
|
|
11
|
+
Parameters
|
|
12
|
+
----------
|
|
13
|
+
num_task : int
|
|
14
|
+
Number of prediction tasks.
|
|
15
|
+
input_dim : int
|
|
16
|
+
Size of vocabulary for SMILES tokenization.
|
|
17
|
+
output_dim : int
|
|
18
|
+
Dimension of embedding vectors.
|
|
19
|
+
LSTMunits : int
|
|
20
|
+
Number of hidden units in LSTM layers.
|
|
21
|
+
max_input_len : int
|
|
22
|
+
Maximum length of input sequences.
|
|
23
|
+
|
|
24
|
+
Attributes
|
|
25
|
+
----------
|
|
26
|
+
embedding : nn.Embedding
|
|
27
|
+
Embedding layer that converts token indices to dense vectors.
|
|
28
|
+
lstm1 : nn.LSTM
|
|
29
|
+
First bidirectional LSTM layer.
|
|
30
|
+
lstm2 : nn.LSTM
|
|
31
|
+
Second bidirectional LSTM layer.
|
|
32
|
+
timedist_dense : nn.Linear
|
|
33
|
+
Time-distributed dense layer for feature transformation.
|
|
34
|
+
relu : nn.ReLU
|
|
35
|
+
ReLU activation function.
|
|
36
|
+
fc : nn.Linear
|
|
37
|
+
Final fully connected layer for prediction.
|
|
38
|
+
|
|
39
|
+
Notes
|
|
40
|
+
-----
|
|
41
|
+
The model architecture consists of:
|
|
42
|
+
1. An embedding layer to convert tokens to vectors
|
|
43
|
+
2. Two stacked bidirectional LSTM layers
|
|
44
|
+
3. A time-distributed dense layer with ReLU activation
|
|
45
|
+
4. A final fully connected layer for prediction
|
|
46
|
+
"""
|
|
47
|
+
|
|
48
|
+
def __init__(self, num_task, input_dim, output_dim, LSTMunits, max_input_len):
|
|
49
|
+
"""
|
|
50
|
+
input_dim: Vocabulary size
|
|
51
|
+
output_dim: Embedding dimension
|
|
52
|
+
LSTMunits: Number of hidden units in LSTM (unidirectional)
|
|
53
|
+
max_input_len: Input sequence length (used for later flattening)
|
|
54
|
+
"""
|
|
55
|
+
super(LSTM, self).__init__()
|
|
56
|
+
self.num_task = num_task
|
|
57
|
+
hidden_dim = int(LSTMunits / 2)
|
|
58
|
+
self.embedding = nn.Embedding(input_dim, output_dim)
|
|
59
|
+
self.lstm1 = nn.LSTM(input_size=output_dim, hidden_size=LSTMunits, bidirectional=True, batch_first=True)
|
|
60
|
+
self.lstm2 = nn.LSTM(input_size=LSTMunits * 2, hidden_size=LSTMunits, bidirectional=True, batch_first=True)
|
|
61
|
+
self.timedist_dense = nn.Linear(LSTMunits * 2, hidden_dim)
|
|
62
|
+
self.relu = nn.ReLU()
|
|
63
|
+
self.fc = nn.Linear(hidden_dim * max_input_len, 1)
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def initialize_parameters(self, seed=None):
|
|
67
|
+
"""Initialize model parameters randomly.
|
|
68
|
+
|
|
69
|
+
Parameters
|
|
70
|
+
----------
|
|
71
|
+
seed : int, optional
|
|
72
|
+
Random seed for reproducibility.
|
|
73
|
+
"""
|
|
74
|
+
if seed is not None:
|
|
75
|
+
torch.manual_seed(seed)
|
|
76
|
+
|
|
77
|
+
# Initialize the main components
|
|
78
|
+
init_weights(self.embedding)
|
|
79
|
+
init_weights(self.lstm1)
|
|
80
|
+
init_weights(self.lstm2)
|
|
81
|
+
init_weights(self.timedist_dense)
|
|
82
|
+
init_weights(self.relu)
|
|
83
|
+
init_weights(self.fc)
|
|
84
|
+
|
|
85
|
+
# Reset all parameters using PyTorch Geometric's reset function
|
|
86
|
+
def reset_parameters(module):
|
|
87
|
+
if hasattr(module, 'reset_parameters'):
|
|
88
|
+
module.reset_parameters()
|
|
89
|
+
elif hasattr(module, 'weight') and hasattr(module.weight, 'data'):
|
|
90
|
+
init_weights(module)
|
|
91
|
+
|
|
92
|
+
self.apply(reset_parameters)
|
|
93
|
+
|
|
94
|
+
def compute_loss(self, batched_input, batched_label, criterion):
|
|
95
|
+
"""Compute the loss for a batch of data.
|
|
96
|
+
|
|
97
|
+
Parameters
|
|
98
|
+
----------
|
|
99
|
+
batched_input : torch.Tensor
|
|
100
|
+
Batch of input sequences, shape (batch_size, seq_len).
|
|
101
|
+
batched_label : torch.Tensor
|
|
102
|
+
Batch of target values, shape (batch_size, 1).
|
|
103
|
+
criterion : callable
|
|
104
|
+
Loss function to use.
|
|
105
|
+
|
|
106
|
+
Returns
|
|
107
|
+
-------
|
|
108
|
+
torch.Tensor
|
|
109
|
+
Scalar loss value.
|
|
110
|
+
"""
|
|
111
|
+
emb = self.embedding(batched_input)
|
|
112
|
+
emb, _ = self.lstm1(emb)
|
|
113
|
+
emb, _ = self.lstm2(emb)
|
|
114
|
+
emb = self.relu(self.timedist_dense(emb))
|
|
115
|
+
emb = emb.contiguous().view(emb.size(0), -1)
|
|
116
|
+
prediction = self.fc(emb)
|
|
117
|
+
target = batched_label.to(torch.float32)
|
|
118
|
+
is_labeled = batched_label == batched_label
|
|
119
|
+
loss = criterion(prediction.to(torch.float32)[is_labeled], target[is_labeled])
|
|
120
|
+
return loss
|
|
121
|
+
|
|
122
|
+
def forward(self, batched_input):
|
|
123
|
+
"""Forward pass of the model.
|
|
124
|
+
|
|
125
|
+
Parameters
|
|
126
|
+
----------
|
|
127
|
+
batched_input : torch.Tensor
|
|
128
|
+
Batch of input sequences, shape (batch_size, seq_len).
|
|
129
|
+
|
|
130
|
+
Returns
|
|
131
|
+
-------
|
|
132
|
+
dict
|
|
133
|
+
Dictionary containing:
|
|
134
|
+
- prediction: Model predictions (shape: [batch_size, 1])
|
|
135
|
+
"""
|
|
136
|
+
# batched_data: (batch_size, seq_len)
|
|
137
|
+
emb = self.embedding(batched_input) # -> (batch, seq_len, output_dim)
|
|
138
|
+
emb, _ = self.lstm1(emb) # -> (batch, seq_len, 2*LSTMunits)
|
|
139
|
+
emb, _ = self.lstm2(emb) # -> (batch, seq_len, 2*LSTMunits)
|
|
140
|
+
emb = self.relu(self.timedist_dense(emb)) # -> (batch, seq_len, hidden_dim)
|
|
141
|
+
emb = emb.contiguous().view(emb.size(0), -1) # flatten: (batch, seq_len * hidden_dim)
|
|
142
|
+
prediction = self.fc(emb) # -> (batch, 1)
|
|
143
|
+
return {"prediction": prediction}
|
|
144
|
+
|
|
145
|
+
|