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.
Files changed (148) hide show
  1. torch_molecule/__init__.py +70 -0
  2. torch_molecule/base/__init__.py +8 -0
  3. torch_molecule/base/base.py +396 -0
  4. torch_molecule/base/encoder.py +44 -0
  5. torch_molecule/base/generator.py +24 -0
  6. torch_molecule/base/old_encoder.py +322 -0
  7. torch_molecule/base/old_predictor.py +904 -0
  8. torch_molecule/base/old_predictor2.py +436 -0
  9. torch_molecule/base/predictor.py +112 -0
  10. torch_molecule/encoder/attrmask/__init__.py +3 -0
  11. torch_molecule/encoder/attrmask/model.py +113 -0
  12. torch_molecule/encoder/attrmask/modeling_attrmask.py +304 -0
  13. torch_molecule/encoder/constant.py +13 -0
  14. torch_molecule/encoder/contextpred/__init__.py +3 -0
  15. torch_molecule/encoder/contextpred/model.py +207 -0
  16. torch_molecule/encoder/contextpred/modeling_contextpred.py +299 -0
  17. torch_molecule/encoder/contextpred/utils.py +206 -0
  18. torch_molecule/encoder/edgepred/__init__.py +3 -0
  19. torch_molecule/encoder/edgepred/model.py +104 -0
  20. torch_molecule/encoder/edgepred/modeling_edgepred.py +285 -0
  21. torch_molecule/encoder/graphmae/__init__.py +3 -0
  22. torch_molecule/encoder/graphmae/dataloader.py +206 -0
  23. torch_molecule/encoder/graphmae/model.py +137 -0
  24. torch_molecule/encoder/graphmae/modeling_graphmae.py +323 -0
  25. torch_molecule/encoder/infograph/__init__.py +3 -0
  26. torch_molecule/encoder/infograph/loss_function.py +128 -0
  27. torch_molecule/encoder/infograph/model.py +92 -0
  28. torch_molecule/encoder/infograph/modeling_infograph.py +301 -0
  29. torch_molecule/encoder/infograph/module.py +31 -0
  30. torch_molecule/encoder/moama/__init__.py +3 -0
  31. torch_molecule/encoder/moama/model.py +131 -0
  32. torch_molecule/encoder/moama/modeling_moama.py +307 -0
  33. torch_molecule/encoder/moama/utils.py +143 -0
  34. torch_molecule/encoder/pretrained/__init__.py +3 -0
  35. torch_molecule/encoder/pretrained/modeling_pretrained.py +274 -0
  36. torch_molecule/encoder/supervised/__init__.py +3 -0
  37. torch_molecule/encoder/supervised/model.py +97 -0
  38. torch_molecule/encoder/supervised/modeling_supervised.py +354 -0
  39. torch_molecule/generator/digress/__init__.py +3 -0
  40. torch_molecule/generator/digress/diffusion.py +194 -0
  41. torch_molecule/generator/digress/modeling_digress.py +503 -0
  42. torch_molecule/generator/digress/transformer.py +360 -0
  43. torch_molecule/generator/digress/utils.py +196 -0
  44. torch_molecule/generator/gdss/__init__.py +3 -0
  45. torch_molecule/generator/gdss/layers.py +257 -0
  46. torch_molecule/generator/gdss/model.py +204 -0
  47. torch_molecule/generator/gdss/modeling_gdss.py +474 -0
  48. torch_molecule/generator/gdss/sde.py +334 -0
  49. torch_molecule/generator/gdss/solver.py +227 -0
  50. torch_molecule/generator/gdss/toremove/ScoreNetworkA.py +167 -0
  51. torch_molecule/generator/gdss/toremove/ScoreNetworkX.py +85 -0
  52. torch_molecule/generator/gdss/utils.py +240 -0
  53. torch_molecule/generator/graph_dit/__init__.py +3 -0
  54. torch_molecule/generator/graph_dit/diffusion.py +188 -0
  55. torch_molecule/generator/graph_dit/modeling_graph_dit.py +581 -0
  56. torch_molecule/generator/graph_dit/transformer.py +207 -0
  57. torch_molecule/generator/graph_dit/utils.py +220 -0
  58. torch_molecule/generator/graph_ga/__init__.py +3 -0
  59. torch_molecule/generator/graph_ga/crossover.py +182 -0
  60. torch_molecule/generator/graph_ga/modeling_graph_ga.py +310 -0
  61. torch_molecule/generator/graph_ga/mutate.py +136 -0
  62. torch_molecule/generator/graph_ga/oracle.py +95 -0
  63. torch_molecule/generator/jtvae/__init__.py +3 -0
  64. torch_molecule/generator/jtvae/jtnn/chemutils.py +450 -0
  65. torch_molecule/generator/jtvae/jtnn/datautils.py +127 -0
  66. torch_molecule/generator/jtvae/jtnn/jtmpn.py +140 -0
  67. torch_molecule/generator/jtvae/jtnn/jtnn_dec.py +350 -0
  68. torch_molecule/generator/jtvae/jtnn/jtnn_enc.py +131 -0
  69. torch_molecule/generator/jtvae/jtnn/mol_tree.py +111 -0
  70. torch_molecule/generator/jtvae/jtnn/mpn.py +129 -0
  71. torch_molecule/generator/jtvae/jtnn/nnutils.py +66 -0
  72. torch_molecule/generator/jtvae/jtnn/vocab.py +31 -0
  73. torch_molecule/generator/jtvae/jtnn_vae.py +246 -0
  74. torch_molecule/generator/jtvae/modeling_jtvae.py +247 -0
  75. torch_molecule/generator/lstm/__init__.py +3 -0
  76. torch_molecule/generator/lstm/action_sampler.py +91 -0
  77. torch_molecule/generator/lstm/lstm.py +59 -0
  78. torch_molecule/generator/lstm/modeling_lstm.py +263 -0
  79. torch_molecule/generator/lstm/smiles_char_dict.py +123 -0
  80. torch_molecule/generator/lstm/utils.py +62 -0
  81. torch_molecule/generator/molgpt/__init__.py +3 -0
  82. torch_molecule/generator/molgpt/dataset.py +87 -0
  83. torch_molecule/generator/molgpt/gpt.py +251 -0
  84. torch_molecule/generator/molgpt/modeling_molgpt.py +427 -0
  85. torch_molecule/generator/pretrained/modeling_pretrained.py +1 -0
  86. torch_molecule/nn/__init__.py +16 -0
  87. torch_molecule/nn/attention.py +63 -0
  88. torch_molecule/nn/embedder.py +183 -0
  89. torch_molecule/nn/gnn.py +631 -0
  90. torch_molecule/nn/mlp.py +34 -0
  91. torch_molecule/predictor/bfgnn/__init__.py +3 -0
  92. torch_molecule/predictor/bfgnn/model.py +116 -0
  93. torch_molecule/predictor/bfgnn/modeling_bfgnn.py +81 -0
  94. torch_molecule/predictor/dir/__init__.py +3 -0
  95. torch_molecule/predictor/dir/model.py +235 -0
  96. torch_molecule/predictor/dir/modeling_dir.py +161 -0
  97. torch_molecule/predictor/dir/utils.py +40 -0
  98. torch_molecule/predictor/gnn/__init__.py +3 -0
  99. torch_molecule/predictor/gnn/model.py +114 -0
  100. torch_molecule/predictor/gnn/modeling_gnn.py +706 -0
  101. torch_molecule/predictor/gnn/old_modeling_gnn.py +707 -0
  102. torch_molecule/predictor/grea/__init__.py +3 -0
  103. torch_molecule/predictor/grea/model.py +190 -0
  104. torch_molecule/predictor/grea/modeling_grea.py +107 -0
  105. torch_molecule/predictor/grea/old1_features.py +229 -0
  106. torch_molecule/predictor/grea/old1_modeling_grea.py +493 -0
  107. torch_molecule/predictor/grea/utils.py +97 -0
  108. torch_molecule/predictor/grin/__init__.py +3 -0
  109. torch_molecule/predictor/grin/model.py +116 -0
  110. torch_molecule/predictor/grin/modeling_grin.py +91 -0
  111. torch_molecule/predictor/irm/__init__.py +3 -0
  112. torch_molecule/predictor/irm/model.py +127 -0
  113. torch_molecule/predictor/irm/modeling_irm.py +179 -0
  114. torch_molecule/predictor/lstm/__init__.py +3 -0
  115. torch_molecule/predictor/lstm/model.py +145 -0
  116. torch_molecule/predictor/lstm/modeling_lstm.py +627 -0
  117. torch_molecule/predictor/lstm/token_from_smiles.py +85 -0
  118. torch_molecule/predictor/old_modeling_graph.py +650 -0
  119. torch_molecule/predictor/rpgnn/__init__.py +3 -0
  120. torch_molecule/predictor/rpgnn/model.py +185 -0
  121. torch_molecule/predictor/rpgnn/modeling_rpgnn.py +76 -0
  122. torch_molecule/predictor/sgir/__init__.py +3 -0
  123. torch_molecule/predictor/sgir/modeling_sgir.py +260 -0
  124. torch_molecule/predictor/sgir/strategy.py +266 -0
  125. torch_molecule/predictor/smiles_transformer/__init__.py +3 -0
  126. torch_molecule/predictor/smiles_transformer/model.py +424 -0
  127. torch_molecule/predictor/smiles_transformer/modeling_transformer.py +252 -0
  128. torch_molecule/predictor/ssr/__init__.py +3 -0
  129. torch_molecule/predictor/ssr/model.py +283 -0
  130. torch_molecule/predictor/ssr/modeling_ssr.py +267 -0
  131. torch_molecule/predictor/ssr/utils.py +185 -0
  132. torch_molecule/utils/__init__.py +31 -0
  133. torch_molecule/utils/checker.py +158 -0
  134. torch_molecule/utils/checkpoint.py +294 -0
  135. torch_molecule/utils/format.py +140 -0
  136. torch_molecule/utils/generic/metrics.py +530 -0
  137. torch_molecule/utils/generic/pseudo_tasks.py +11 -0
  138. torch_molecule/utils/generic/weights.py +30 -0
  139. torch_molecule/utils/graph/features.py +178 -0
  140. torch_molecule/utils/graph/graph_from_smiles.py +134 -0
  141. torch_molecule/utils/graph/graph_to_smiles.py +312 -0
  142. torch_molecule/utils/hf.py +232 -0
  143. torch_molecule/utils/search.py +64 -0
  144. torch_molecule-0.1.0.dev2.dist-info/METADATA +243 -0
  145. torch_molecule-0.1.0.dev2.dist-info/RECORD +148 -0
  146. torch_molecule-0.1.0.dev2.dist-info/WHEEL +5 -0
  147. torch_molecule-0.1.0.dev2.dist-info/licenses/LICENSE +21 -0
  148. 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,3 @@
1
+ from .modeling_irm import IRMMolecularPredictor
2
+
3
+ __all__ = ['IRMMolecularPredictor']
@@ -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,3 @@
1
+ from .modeling_lstm import LSTMMolecularPredictor
2
+
3
+ __all__ = ['LSTMMolecularPredictor']
@@ -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
+