D4CMPP2 0.4.0__py3-none-any.whl

D4CMPP2/networks/registry.py

@@ -0,0 +1,187 @@
+"""Typed network registry used by training and optimization."""
+
+from dataclasses import dataclass
+
+from D4CMPP2.networks.base import MolecularNetwork
+
+
+@dataclass(frozen=True)
+class ModelDefinition:
+    name: str
+    network: type[MolecularNetwork]
+    aliases: tuple[str, ...] = ()
+    data_contract: str = "molecule"
+
+    def __post_init__(self):
+        if not isinstance(self.network, type) or not issubclass(
+            self.network, MolecularNetwork
+        ):
+            raise TypeError(
+                f"Registered network for {self.name!r} must inherit MolecularNetwork."
+            )
+        if "<locals>" in self.network.__qualname__:
+            raise ValueError(
+                f"Network {self.network.__qualname__!r} must be defined at module "
+                "scope so saved runs can import it."
+            )
+        if self.network.model_name != self.name:
+            raise ValueError(
+                f"Registry name {self.name!r} does not match "
+                f"{self.network.__name__}.model_name={self.network.model_name!r}."
+            )
+        self.network.optimization_space()
+        if self.data_contract not in {"molecule", "solvent", "isa"}:
+            raise ValueError(
+                f"Model {self.name!r} has unknown data_contract "
+                f"{self.data_contract!r}. Expected 'molecule', 'solvent', or 'isa'."
+            )
+
+    def training_config(self):
+        managers = {
+            "molecule": ("MolDataManager", "MolDataManager", "NetworkManager", "NetworkManager", "TrainManager", "Trainer"),
+            "solvent": ("MolDataManager", "MolDataManager_withSolv", "NetworkManager", "NetworkManager", "TrainManager", "Trainer"),
+            "isa": ("ISADataManager", "ISADataManager", "ISANetworkManager", "ISANetworkManager", "ISATrainManager", "ISATrainer"),
+        }
+        data_module, data_class, network_module, network_class, train_module, train_class = managers[self.data_contract]
+        return {
+            "name": self.name,
+            "network": self.network.__module__.rsplit(".", 1)[-1],
+            "network_id": self.name,
+            "data_manager_module": data_module,
+            "data_manager_class": data_class,
+            "network_manager_module": network_module,
+            "network_manager_class": network_class,
+            "train_manager_module": train_module,
+            "train_manager_class": train_class,
+            "description": self.network.__doc__ or self.name,
+            "version": "2.0",
+        }
+
+
+_MODELS = {}
+
+
+def register_model(definition):
+    names = (definition.name, *definition.aliases)
+    conflicts = [name for name in names if name in _MODELS]
+    if conflicts:
+        raise ValueError(f"Model names are already registered: {conflicts!r}.")
+    for name in names:
+        _MODELS[name] = definition
+    return definition
+
+
+def register_network(network, *, aliases=(), data_contract="molecule"):
+    """Register a MolecularNetwork subclass for train and optimize entry points."""
+    return register_model(
+        ModelDefinition(
+            network.model_name,
+            network,
+            aliases=tuple(aliases),
+            data_contract=data_contract,
+        )
+    )
+
+
+def get_model(name):
+    try:
+        return _MODELS[name]
+    except KeyError as exc:
+        raise ValueError(
+            f"Unknown model {name!r}. Registered names: {sorted(_MODELS)!r}."
+        ) from exc
+
+
+def registered_models():
+    return {
+        definition.name: definition
+        for definition in _MODELS.values()
+    }
+
+
+def _register_builtins():
+    from D4CMPP2.networks.AFP_model import AFP
+    from D4CMPP2.networks.AFPwithSolv_model import SolventAFP
+    from D4CMPP2.networks.DMPNN_model import DMPNN
+    from D4CMPP2.networks.DMPNNwithSolv_model import SolventDMPNN
+    from D4CMPP2.networks.GAT_model import GAT
+    from D4CMPP2.networks.GATwithSolv_model import SolventGAT
+    from D4CMPP2.networks.GCN_model import GCN
+    from D4CMPP2.networks.GCNwithSolv_model import SolventGCN
+    from D4CMPP2.networks.GC_model import GroupContributionNetwork
+    from D4CMPP2.networks.ISAT_model import ISAT
+    from D4CMPP2.networks.ISATPN_model import ISATPN
+    from D4CMPP2.networks.MPNN_model import MPNN
+    from D4CMPP2.networks.MPNNwithSolv_model import SolventMPNN
+
+    register_model(ModelDefinition("afp", AFP, aliases=("AFP", "AFP_model")))
+    register_model(
+        ModelDefinition(
+            "afp_solvent",
+            SolventAFP,
+            aliases=("AFPwS", "AFPwithSolv_model"),
+            data_contract="solvent",
+        )
+    )
+    register_model(
+        ModelDefinition("dmpnn", DMPNN, aliases=("DMPNN", "DMPNN_model"))
+    )
+    register_model(
+        ModelDefinition(
+            "dmpnn_solvent",
+            SolventDMPNN,
+            aliases=("DMPNNwS", "DMPNNwithSolv_model"),
+            data_contract="solvent",
+        )
+    )
+    register_model(ModelDefinition("gcn", GCN, aliases=("GCN", "GCN_model")))
+    register_model(
+        ModelDefinition(
+            "gcn_solvent",
+            SolventGCN,
+            aliases=("GCNwS", "GCNwithSolv_model"),
+            data_contract="solvent",
+        )
+    )
+    register_model(
+        ModelDefinition(
+            "group_contribution",
+            GroupContributionNetwork,
+            aliases=("GC", "GC_model"),
+            data_contract="isa",
+        )
+    )
+    register_model(
+        ModelDefinition(
+            "isat", ISAT, aliases=("ISAT", "ISAT_model"), data_contract="isa"
+        )
+    )
+    register_model(
+        ModelDefinition(
+            "isatpn",
+            ISATPN,
+            aliases=("ISATPN", "ISATPM", "ISATPM_model"),
+            data_contract="isa",
+        )
+    )
+    register_model(ModelDefinition("gat", GAT, aliases=("GAT", "GAT_model")))
+    register_model(
+        ModelDefinition(
+            "gat_solvent",
+            SolventGAT,
+            aliases=("GATwS", "GATwithSolv_model"),
+            data_contract="solvent",
+        )
+    )
+    register_model(ModelDefinition("mpnn", MPNN, aliases=("MPNN", "MPNN_model")))
+    register_model(
+        ModelDefinition(
+            "mpnn_solvent",
+            SolventMPNN,
+            aliases=("MPNNwS", "MPNNwithSolv_model"),
+            data_contract="solvent",
+        )
+    )
+
+
+_register_builtins()

D4CMPP2/networks/src/AFP.py

@@ -0,0 +1,118 @@
+
+"""
+This codes are modified from the project "GC-GNN" (https://github.com/gsi-lab/GC-GNN)
+The original codes are under the MIT License. (https://github.com/gsi-lab/GC-GNN/blob/main/networks/AttentiveFP.py)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from D4CMPP2.networks.src.GCN import graph_sum_pool
+
+
+def _graph_softmax(values, batch, num_graphs):
+    maximum = values.new_full((num_graphs, values.shape[1]), -torch.inf)
+    maximum.scatter_reduce_(0, batch[:, None].expand_as(values), values, reduce='amax', include_self=True)
+    exponentials = torch.exp(values - maximum[batch])
+    denominator = values.new_zeros((num_graphs, values.shape[1]))
+    denominator.index_add_(0, batch, exponentials)
+    return exponentials / denominator[batch]
+
+class Atom_AFPLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_dim = config['hidden_dim']
+        self.embedding_edge_lin = nn.Sequential(
+            nn.Linear(self.hidden_dim + self.hidden_dim, self.hidden_dim, bias=True),
+            nn.Dropout(config['dropout']),
+            nn.LeakyReLU()
+        )
+        self.cal_alignment = nn.Sequential(
+            nn.Linear(self.hidden_dim + self.hidden_dim, 1, bias=True),
+        )
+        self.attend = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.hidden_dim, bias=True),
+            nn.Dropout(config['dropout']),
+        )
+        self.GRUCell = nn.GRUCell(self.hidden_dim, self.hidden_dim)
+        self.dropout = nn.Dropout(config['dropout'])
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.embedding_edge_lin[0].reset_parameters()
+        self.cal_alignment[0].reset_parameters()
+        self.attend[0].reset_parameters()
+        self.GRUCell.reset_parameters()
+
+    def forward(self, graph, node, edge):
+        src, dst = graph.edge_index
+        neighbor_message = self.embedding_edge_lin(torch.cat([node[src], edge], dim=-1))
+        # Retained for state-dict and numerical compatibility although the legacy path did not apply the score.
+        self.cal_alignment(torch.cat([neighbor_message, node[dst]], dim=-1))
+        messages = self.dropout(self.attend(neighbor_message))
+        context = messages.new_zeros(node.shape)
+        context.index_add_(0, dst, messages)
+        new_node = node + context
+        return new_node
+
+class Mol_AFPLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.GRUCell = nn.GRUCell(config['hidden_dim'], config['hidden_dim'])
+
+        self.cal_alignment = nn.Sequential(
+            nn.Linear(config['hidden_dim'] + config['hidden_dim'], 1, bias=True),
+            nn.LeakyReLU()
+        )
+        self.attend = nn.Sequential(
+            nn.Linear(config['hidden_dim'], config['hidden_dim'], bias=True),
+            nn.Dropout(config['dropout']),
+        )
+        self.dropout = nn.Dropout(config['dropout'])
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.GRUCell.reset_parameters()
+        self.cal_alignment[0].reset_parameters()
+        self.attend[0].reset_parameters()
+
+    def forward(self, graph, super_node, node):
+        super_node = F.leaky_relu(super_node)
+        batch = getattr(graph, 'batch', node.new_zeros(node.shape[0], dtype=torch.long))
+        num_graphs = int(getattr(graph, 'num_graphs', 1))
+        score = self.cal_alignment(torch.cat([node, super_node[batch]], dim=1))
+        attention_weight = _graph_softmax(score, batch, num_graphs)
+        hidden_node = self.attend(self.dropout(node)) * attention_weight
+        super_context = hidden_node.new_zeros((num_graphs, hidden_node.shape[1]))
+        super_context.index_add_(0, batch, hidden_node)
+        super_context = F.elu(super_context)
+        
+        super_node = F.relu(self.GRUCell(super_node, super_context))
+        return super_node, attention_weight
+
+
+class AttentiveFP(nn.Module):
+    # Generate Context of each nodes
+    def __init__(self, config):
+        super().__init__()
+        conv_layers = config.get('conv_layers', 4)
+        T = config.get("T", conv_layers)
+        self.PassingDepth = nn.ModuleList([Atom_AFPLayer(config) for _ in range(conv_layers)])
+        self.MultiTimeSteps = nn.ModuleList([Mol_AFPLayer(config) for d in range(T)])
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for l in self.MultiTimeSteps:
+            l.reset_parameters()
+
+    def forward(self, graph, node, edge, **kwargs):
+        for i in range(len(self.PassingDepth)):
+            node = self.PassingDepth[i](graph, node, edge)
+        if kwargs.get('only_atom', False):
+            return node, None
+        super_node = graph_sum_pool(graph, node)
+        for i in range(len(self.MultiTimeSteps)):
+            super_node, att_w = self.MultiTimeSteps[i](graph, super_node, node)
+        return super_node, att_w
+
+

D4CMPP2/networks/src/BiDropout.py

@@ -0,0 +1,29 @@
+import torch
+import torch.nn as nn
+
+class Bi_Dropout(nn.Module):
+    def __init__(self, p):
+        super(Bi_Dropout, self).__init__()
+        self.p = p
+        self.drop_p = p
+        self.training = True
+
+    def set_drop_p(self,trainig):
+        self.training = trainig
+        self.drop_p = torch.rand(1).item()
+
+    def drop_label(self, score:torch.Tensor,target:torch.Tensor):
+        if self.training:    
+            if self.drop_p<self.p: 
+                target=nn.ReLU()(target)
+            elif self.drop_p>1-self.p: 
+                target=-nn.ReLU()(-target)
+        return score,target
+    
+    def forward(self, x):
+        if self.training:
+            if self.drop_p<self.p: 
+                x[:,1]= 0
+            elif self.drop_p>1-self.p:
+                x[:,0]= 0
+        return x

D4CMPP2/networks/src/DMPNN.py

@@ -0,0 +1,35 @@
+
+"""
+This codes are modified from the project "GC-GNN" (https://github.com/gsi-lab/GC-GNN)
+The original codes are under the MIT License. (https://github.com/gsi-lab/GC-GNN/blob/main/networks/DMPNN.py)
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DMPNNLayer(nn.Module):
+    def __init__(self, in_feats, edge_feats, out_feats, activation, dropout=0.2):
+        super(DMPNNLayer, self).__init__()
+        self.activation = activation
+
+        self.W_m = nn.Linear(in_feats+edge_feats, out_feats, bias=True)
+        self.dropout_layer = nn.Dropout(dropout)
+
+    
+    def forward(self, graph, node_feats, edge_feats, direct_feats=None, backward_feats=None):
+        src, dst = graph.edge_index
+        if backward_feats is not None and direct_feats is not None:
+            direct = self.W_m(torch.cat([edge_feats, direct_feats], 1))
+            backward = self.W_m(torch.cat([edge_feats, backward_feats], 1))
+        else:
+            direct = self.W_m(torch.cat([edge_feats, node_feats[src]], 1))
+            backward = self.W_m(torch.cat([edge_feats, node_feats[dst]], 1))
+        full_feats = direct.new_zeros((node_feats.shape[0], direct.shape[1]))
+        full_feats.index_add_(0, dst, direct)
+        new_edge_feats = full_feats[src] - backward
+        new_backward_feats = full_feats[dst] - direct
+        new_node_feats = full_feats
+        if self.activation is not None:
+            new_node_feats = self.activation(new_node_feats)
+        new_node_feats = self.dropout_layer(new_node_feats)
+        return new_node_feats, new_edge_feats, new_backward_feats

D4CMPP2/networks/src/GAT.py

@@ -0,0 +1,69 @@
+import torch.nn as nn
+import torch
+
+class GAT_layer(nn.Module):
+    def __init__(self, in_node_feats, hidden_feats, out_feats, activation, dropout=0.2, batch_norm=False, residual_sum = False):
+        super(GAT_layer, self).__init__()
+        self.activation = activation
+        self.dropout = nn.Dropout(dropout)
+
+        self.attention_W = nn.Linear(in_node_feats*2, hidden_feats)
+        self.attention_a = nn.Linear(hidden_feats,1, bias=False)
+        self.linear = nn.Linear(in_node_feats, out_feats)
+
+        self.batch_norm = batch_norm
+        self.residual_sum = residual_sum
+        if self.batch_norm:
+            self.bn = nn.BatchNorm1d(out_feats)
+        if residual_sum:
+            if in_node_feats!=out_feats:
+                self.residual_layer = nn.Linear(in_node_feats, out_feats)
+
+    def forward(self, graph, node_feats):
+        src, dst = graph.edge_index
+        logits = self.attention_a(
+            nn.LeakyReLU()(self.attention_W(torch.cat([node_feats[src], node_feats[dst]], dim=1)))
+        )
+        max_per_dst = logits.new_full((node_feats.shape[0], logits.shape[1]), -torch.inf)
+        max_per_dst.scatter_reduce_(
+            0, dst[:, None].expand_as(logits), logits, reduce='amax', include_self=True
+        )
+        exponentials = torch.exp(logits - max_per_dst[dst])
+        denominator = logits.new_zeros((node_feats.shape[0], logits.shape[1]))
+        denominator.index_add_(0, dst, exponentials)
+        scores = exponentials / denominator[dst]
+
+        messages = scores * self.linear(node_feats[src])
+        h = messages.new_zeros((node_feats.shape[0], messages.shape[1]))
+        h.index_add_(0, dst, messages)
+        h = nn.LeakyReLU()(h)
+        if self.batch_norm:
+            h = self.bn(h)
+        h = self.activation(h)
+        h = self.dropout(h)
+        if self.residual_sum:
+            if node_feats.shape[1]!=h.shape[1]:
+                node_feats = self.residual_layer(node_feats)
+            h = h + node_feats
+        return h
+    
+class GATs(nn.Module):
+    def __init__(self,  in_node_feats, hidden_feats, out_feats, activation, n_layers, dropout=0.2, batch_norm=False, residual_sum = False):
+        super(GATs, self).__init__()
+
+        self.layers = nn.ModuleList()
+        for i in range(n_layers):
+            if i==0:
+                _in_feats = in_node_feats
+            else:
+                _in_feats = hidden_feats
+            if i==n_layers-1:
+                _out_feats = out_feats
+            else:
+                _out_feats = hidden_feats
+            self.layers.append(GAT_layer(_in_feats, hidden_feats, _out_feats, activation, dropout, batch_norm, residual_sum))
+        
+    def forward(self, graph, node_feats):
+        for layer in self.layers:
+            node_feats = layer(graph, node_feats)
+        return node_feats

D4CMPP2/networks/src/GC.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from .GCN import GCN_layer 
+from .MPNN import MPNN_layer
+from .distGCN import distGCN_layer
+from .pyg_hetero import relation_graph, relation_sum
+
+class GCconvolution(nn.Module):
+    def __init__(self, in_node_feats, in_edge_feats, out_feats, activation, n_layers, dropout=0.2, batch_norm=False, residual_sum = False, alpha=0.1, max_dist = 4):
+        super().__init__()        
+        # Message Passing
+        self.r2r = nn.ModuleList([MPNN_layer(in_node_feats, in_edge_feats, out_feats, activation, dropout, batch_norm, residual_sum) for _ in range(n_layers)])
+        self.i2i = nn.ModuleList([GCN_layer(out_feats, out_feats, activation, dropout, batch_norm, residual_sum) for _ in range(n_layers)])
+
+        self.r2i = r2i_layer()
+        self.i2d = i2s_layer()
+        self.d2d = distGCN_layer(out_feats, max_dist, out_feats, activation, alpha)
+
+        self.d2score = nn.Sequential(nn.Linear(out_feats, out_feats//2), 
+                                     nn.LeakyReLU(),
+                                     nn.Linear(out_feats//2, out_feats//4),
+                                     nn.LeakyReLU(),
+                                     nn.Linear(out_feats//4, 1)
+                                     )
+    
+    def forward(self, graph, r_node, r2r_edge, i_node, d2d_edge):
+                
+        real_graph = relation_graph(graph, 'r_nd', 'r2r')
+        image_graph = relation_graph(graph, 'i_nd', 'i2i')
+
+        for i in range(len(self.r2r)):
+            r_node = self.r2r[i](real_graph, r_node, r2r_edge)
+            i_node = self.r2i(graph, r_node, i_node)
+            i_node = self.i2i[i](image_graph, i_node)
+        d_node = self.i2d(graph, i_node)
+        score= self.d2score(d_node)
+        
+        return score
+        
+class GCconvolution2(nn.Module):
+    def __init__(self, in_node_feats, in_edge_feats, out_feats, activation, n_layers, dropout=0.2, batch_norm=False, residual_sum = False, alpha=0.1, max_dist = 4):
+        super().__init__()        
+        # Message Passing
+        self.r2r = nn.ModuleList([GCN_layer(out_feats, out_feats, activation, dropout, batch_norm, residual_sum) for _ in range(n_layers)])
+        self.i2i = nn.ModuleList([GCN_layer(out_feats, out_feats, activation, dropout, batch_norm, residual_sum) for _ in range(n_layers)])
+
+        self.r2i = r2i_layer()
+        self.i2d = i2s_layer()
+        self.d2d = distGCN_layer(out_feats, max_dist, out_feats, activation, alpha)
+
+        self.d2score = nn.Sequential(nn.Linear(out_feats, out_feats//2), 
+                                     nn.LeakyReLU(),
+                                     nn.Linear(out_feats//2, out_feats//4),
+                                     nn.LeakyReLU(),
+                                     nn.Linear(out_feats//4, 1)
+                                     )
+    
+    def forward(self, graph, r_node, r2r_edge, i_node, d2d_edge):
+                
+        image_graph = relation_graph(graph, 'i_nd', 'i2i')
+
+        for i in range(len(self.r2r)):
+            r_node = self.r2r[i](image_graph, r_node)
+            i_node = self.r2i(graph, r_node, i_node)
+            i_node = self.i2i[i](image_graph, i_node)
+        d_node = self.i2d(graph, i_node)
+        score= self.d2score(d_node)
+        
+        return score
+        
+class r2i_layer(nn.Module):
+    def forward(self,graph, r_node, i_node):
+        return i_node+r_node
+    
+class i2s_layer(nn.Module):
+    def forward(self,graph, i_node):
+        return relation_sum(graph, 'i_nd', 'i2d', 'd_nd', i_node)
+
+
+class s2r_Layer(nn.Module):    
+    def forward(self, graph, node):
+        return relation_sum(graph, 'd_nd', 'd2r', 'r_nd', node)

D4CMPP2/networks/src/GCN.py

@@ -0,0 +1,71 @@
+import torch.nn as nn
+import torch
+
+
+class GCN_layer(nn.Module):
+    def __init__(self, in_feats, out_feats, activation, dropout=0.2, batch_norm=False, residual_sum = False):
+        super(GCN_layer, self).__init__()
+        self.activation = activation
+        self.dropout = nn.Dropout(dropout)
+        self.batch_norm = batch_norm
+        self.linear = nn.Linear(in_feats, out_feats)
+        self.residual_sum = residual_sum
+        if self.batch_norm:
+            self.bn = nn.BatchNorm1d(out_feats)
+        if residual_sum:
+            if in_feats!=out_feats:
+                self.residual_layer = nn.Linear(in_feats, out_feats)
+
+    def forward(self, graph, node_feats):
+        h = self.linear(node_feats)
+        src, dst = graph.edge_index
+        aggregated = torch.zeros_like(h)
+        if src.numel() > 0:
+            aggregated.index_add_(0, dst, h[src])
+        # The established contract adds one self-message per node without removing
+        # pre-existing self edges. Existing self edges are already aggregated above.
+        h = aggregated + h
+        if self.batch_norm:
+            h = self.bn(h)
+        h = self.activation(h)
+        h = self.dropout(h)
+        if self.residual_sum:
+            if node_feats.shape[1]!=h.shape[1]:
+                node_feats = self.residual_layer(node_feats)
+            h = h + node_feats
+        return h
+    
+class GCNs(nn.Module):
+    def __init__(self,  in_feats, hidden_feats, out_feats, activation, n_layers, dropout=0.2, batch_norm=False, residual_sum = False):
+        super(GCNs, self).__init__()
+
+        self.layers = nn.ModuleList()
+        for i in range(n_layers):
+            if i==0:
+                _in_feats = in_feats
+            else:
+                _in_feats = hidden_feats
+            if i==n_layers-1:
+                _out_feats = out_feats
+            else:
+                _out_feats = hidden_feats
+            self.layers.append(GCN_layer(_in_feats, _out_feats, activation, dropout, batch_norm, residual_sum))
+        
+    def forward(self, graph, node_feats):
+        for layer in self.layers:
+            node_feats = layer(graph, node_feats)
+        return node_feats
+
+
+def graph_sum_pool(graph, node_feats):
+    """Sum node features per PyG graph using the established readout contract."""
+    if hasattr(graph, "batch") and graph.batch is not None:
+        batch = graph.batch
+        num_graphs = int(graph.num_graphs)
+    else:
+        batch = torch.zeros(node_feats.shape[0], dtype=torch.long, device=node_feats.device)
+        num_graphs = 1
+    pooled = node_feats.new_zeros((num_graphs, node_feats.shape[-1]))
+    if batch.numel() > 0:
+        pooled.index_add_(0, batch, node_feats)
+    return pooled