D4CMPP2 0.4.0__py3-none-any.whl

D4CMPP2/networks/src/ISAT.py

@@ -0,0 +1,153 @@
+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
+
+
+def _unbiased_variance_or_zero(values):
+    """Preserve unbiased variance except for its undefined one-value case."""
+    flattened = values.view(-1)
+    if flattened.numel() == 1:
+        return flattened.sum() * 0.0
+    return torch.var(flattened)
+
+
+class ISATconvolution(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, 1)
+                                     )
+        self.d2r = s2r_Layer()
+        self.batch_norm = nn.BatchNorm1d(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')
+        dot_graph = relation_graph(graph, 'd_nd', 'd2d')
+
+        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)
+        d_node = self.d2d(dot_graph, d_node, d2d_edge)
+        score= self.d2r(graph, d_node)
+        score= self.d2score(score)
+        score= nn.Sigmoid()(self.batch_norm(score))
+        
+        r_node=score*r_node
+        return r_node, score
+        
+class ISATconvolution_PM(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_P = nn.ModuleList([GCN_layer(out_feats, out_feats, activation, dropout, batch_norm, residual_sum) for _ in range(n_layers)])
+        self.d2d_P = distGCN_layer(out_feats, max_dist, out_feats, activation, alpha)
+        self.d2score_P = nn.Sequential(nn.Linear(out_feats, out_feats//2), 
+                                     nn.LeakyReLU(),
+                                     nn.Linear(out_feats//2, 3),
+                                     nn.LeakyReLU()
+                                     )
+        self.d2score_pl = nn.Linear(3, 1)
+
+        self.i2i_M = nn.ModuleList([GCN_layer(out_feats, out_feats, activation, dropout, batch_norm, residual_sum) for _ in range(n_layers)])
+        self.d2d_M = distGCN_layer(out_feats, max_dist, out_feats, activation, alpha)
+        self.d2score_M = nn.Sequential(nn.Linear(out_feats, out_feats//2), 
+                                     nn.LeakyReLU(),
+                                     nn.Linear(out_feats//2,3),
+                                     nn.LeakyReLU()
+                                     )
+        self.d2score_ml = nn.Linear(3, 1)
+        
+        self.r2i = r2i_layer()
+        self.i2d = i2s_layer()
+        self.d2r = s2r_Layer()
+
+        self.batch_norm_P = nn.BatchNorm1d(1)
+        self.batch_norm_M = nn.BatchNorm1d(1)
+        
+    
+    def forward(self, graph, r_node, r2r_edge, i_node, d2d_edge, **kargs):
+                
+        real_graph = relation_graph(graph, 'r_nd', 'r2r')
+        image_graph = relation_graph(graph, 'i_nd', 'i2i')
+        dot_graph = relation_graph(graph, 'd_nd', 'd2d')
+
+        i_node_P = i_node
+        i_node_M = i_node
+        for i in range(len(self.r2r)):
+            r_node = self.r2r[i](real_graph, r_node, r2r_edge)
+
+            i_node_P = self.r2i(graph, r_node, i_node_P)
+            i_node_P = self.i2i_P[i](image_graph, i_node_P)
+
+            i_node_M = self.r2i(graph, r_node, i_node_M)
+            i_node_M = self.i2i_M[i](image_graph, i_node_M)
+
+
+        d_node_P = self.i2d(graph, i_node_P)
+        score_P = self.d2d_P(dot_graph, d_node_P, d2d_edge)
+        score_P_temp= score_P.clone()
+        score_P= self.d2score_P(score_P)
+        score_P= self.d2score_pl(score_P)
+        score_P= nn.Sigmoid()(self.batch_norm_P(score_P))
+        score_P_node= self.d2r(graph, score_P)
+
+        r_node_P=score_P_node*r_node
+
+
+        d_node_M = self.i2d(graph, i_node_M)
+        score_M = self.d2d_M(dot_graph, d_node_M, d2d_edge)     
+        score_M_temp= score_M.clone()
+        score_M= self.d2score_M(score_M)
+        score_M= self.d2score_ml(score_M)   
+        score_M= nn.Sigmoid()(self.batch_norm_M(score_M))
+        score_M_node= self.d2r(graph, score_M)
+
+        r_node_M=score_M_node*r_node
+
+        if kargs.get('get_feature',False):
+            return score_P_temp, score_M_temp, score_P, score_M # score by group, score by atom
+        
+        self.p_score_var = _unbiased_variance_or_zero(score_P)
+        self.n_score_var = _unbiased_variance_or_zero(score_M)
+        self.p_score_mean = torch.mean(score_P.view(-1))
+        self.n_score_mean = torch.mean(score_M.view(-1))
+        self.p_score_ms = torch.mean(torch.square(0.5-score_P.view(-1)))
+        self.n_score_ms = torch.mean(torch.square(0.5-score_M.view(-1)))
+        return r_node_P, r_node_M, score_P, score_M
+
+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/Linear.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+
+class Linears(nn.Module):
+    def __init__(self, in_feats, out_feats, activation, n_layers, dropout=0.2, batch_norm=False, residual_sum = False, last= False):
+        super().__init__()
+        self.activation = activation
+        self.dropout = nn.Dropout(dropout)
+        self.batch_norm = batch_norm
+        self.linears = nn.ModuleList()
+        self.batch_norms = nn.ModuleList()
+        self.residual_sum = residual_sum
+        self.last = last
+        hidden_dim = (in_feats+out_feats)//2
+        for i in range(n_layers):
+            if i==0:
+                _in_feats = in_feats
+            else:
+                _in_feats = hidden_dim
+            if i==n_layers-1:
+                _out_feats = out_feats
+            else:
+                _out_feats = hidden_dim
+            self.linears.append(nn.Linear(_in_feats, _out_feats))
+            if self.batch_norm:
+                self.batch_norms.append(nn.BatchNorm1d(_out_feats))
+            if residual_sum:
+                if in_feats!=out_feats:
+                    self.residual_layer = nn.Linear(_in_feats, _out_feats) 
+
+    def forward(self, x):
+        count=0
+        for linear in self.linears:
+            count+=1
+            h = linear(x)
+            if count==len(self.linears) and self.last:
+                return h
+            else:
+                if self.batch_norm:
+                    h = self.batch_norms[count-1](h)
+                h = self.activation(h)
+            h = self.dropout(h)
+            if self.residual_sum:
+                if x.shape[1]==h.shape[1]:
+                    x = h + x
+            else:
+                x = h
+        return x
+    

D4CMPP2/networks/src/MPNN.py

@@ -0,0 +1,56 @@
+import torch.nn as nn
+import torch
+
+
+class MPNN_layer(nn.Module):
+    def __init__(self, in_node_feats, in_edge_feats, out_feats, activation, dropout=0.2, batch_norm=False, residual_sum = False):
+        super(MPNN_layer, self).__init__()
+        self.activation = activation
+        self.dropout = nn.Dropout(dropout)
+        self.batch_norm = batch_norm
+        self.linear = nn.Linear(in_node_feats*2+in_edge_feats, out_feats)
+        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, edge_feats):
+        src, dst = graph.edge_index
+        messages = self.linear(torch.cat([node_feats[src], node_feats[dst], edge_feats], dim=1))
+        h = messages.new_zeros((node_feats.shape[0], messages.shape[1]))
+        h.index_add_(0, dst, messages)
+        degree = messages.new_zeros((node_feats.shape[0], 1))
+        degree.index_add_(0, dst, messages.new_ones((messages.shape[0], 1)))
+        h = h / degree.clamp_min(1)
+        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 MPNNs(nn.Module):
+    def __init__(self,  in_node_feats, in_edge_feats, hidden_feats, out_feats, activation, n_layers, dropout=0.2, batch_norm=False, residual_sum = False):
+        super(MPNNs, 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(MPNN_layer(_in_feats, in_edge_feats, _out_feats, activation, dropout, batch_norm, residual_sum))
+        
+    def forward(self, graph, node_feats, edge_feats):
+        for layer in self.layers:
+            node_feats = layer(graph, node_feats, edge_feats)
+        return node_feats

D4CMPP2/networks/src/SolventLayer.py

@@ -0,0 +1,62 @@
+
+from D4CMPP2.networks.src.GCN import GCNs, graph_sum_pool
+from D4CMPP2.networks.src.Linear import Linears
+import torch.nn as nn
+import torch
+
+class SolventLayer(nn.Module):
+    def __init__(self, config):
+        super(SolventLayer, self).__init__()
+        
+        hidden_dim = config["hidden_dim"]
+        linear_layers = config["linear_layers"]
+        solvent_hidden_dim = config.get("solvent_hidden_dim", 64)
+        solvent_conv_layers = config.get("solvent_conv_layers", 4)
+        solvent_linear_layers = config.get("solvent_linear_layers", 2)
+        solvent_dropout = config.get("solvent_dropout", 0.2)
+        self.node_embedding = nn.Linear(config['node_dim'], hidden_dim)
+        self.node_embedding_solv = nn.Linear(
+            config["node_dim"], solvent_hidden_dim
+        )
+
+        self.GCNs_solv = GCNs(
+            solvent_hidden_dim,
+            solvent_hidden_dim,
+            solvent_hidden_dim,
+            nn.ReLU(),
+            solvent_conv_layers,
+            solvent_dropout,
+            False,
+            True,
+        )
+        self.Linears2 = Linears(
+            solvent_hidden_dim,
+            solvent_hidden_dim,
+            nn.ReLU(),
+            solvent_linear_layers,
+            solvent_dropout,
+            False,
+            False,
+        )
+        self.Linears3 = Linears(
+            hidden_dim + solvent_hidden_dim,
+            config["target_dim"],
+            nn.ReLU(),
+            linear_layers,
+            config.get("dropout", 0.2),
+            False,
+            False,
+            True,
+        )
+
+    def forward(self, hidden_feats,solv_graph,solv_node_feats,**kwargs):
+        h = hidden_feats
+
+        h_solv = self.node_embedding_solv(solv_node_feats)
+        h_solv = self.GCNs_solv(solv_graph, h_solv)
+        h_solv = graph_sum_pool(solv_graph, h_solv)
+
+        h_solv = self.Linears2(h_solv)
+        h = torch.cat([h,h_solv],axis=1)
+        h = self.Linears3(h)
+        return h

D4CMPP2/networks/src/distGCN.py

@@ -0,0 +1,21 @@
+import torch.nn as nn
+
+
+class distGCN_layer(nn.Module):
+    def __init__(self, in_feats, edge_feats, out_feats, activation, alpha=0.1):
+        super().__init__()
+        self.activation = activation
+        self.linear = nn.Linear(in_feats, out_feats)
+        self.edge_linear = nn.Linear(edge_feats, out_feats)
+        self.alpha = alpha
+
+    def forward(self, graph, node_feats, edge_feats):
+        if len(edge_feats) == 0:
+            return node_feats
+        src, dst = graph.edge_index
+        transformed_nodes = self.linear(node_feats)
+        transformed_edges = self.edge_linear(edge_feats)
+        messages = transformed_nodes[src] * transformed_edges
+        aggregated = messages.new_zeros((node_feats.shape[0], messages.shape[1]))
+        aggregated.index_add_(0, dst, messages)
+        return node_feats + self.alpha * self.activation(aggregated)

D4CMPP2/networks/src/pyg_hetero.py

@@ -0,0 +1,24 @@
+from types import SimpleNamespace
+
+import torch
+
+
+def relation_graph(graph, node_type, relation):
+    """Expose one homogeneous HeteroData relation to the existing PyG layers."""
+    store = graph[node_type]
+    batch = getattr(store, 'batch', None)
+    if batch is None:
+        batch = torch.zeros(store.num_nodes, dtype=torch.long, device=store.x.device)
+    return SimpleNamespace(
+        edge_index=graph[node_type, relation, node_type].edge_index,
+        batch=batch,
+        num_graphs=int(getattr(graph, 'num_graphs', 1)),
+    )
+
+
+def relation_sum(graph, source_type, relation, destination_type, source_features):
+    """Sum source features at each destination node for one typed relation."""
+    src, dst = graph[source_type, relation, destination_type].edge_index
+    output = source_features.new_zeros((graph[destination_type].num_nodes, source_features.shape[1]))
+    output.index_add_(0, dst, source_features[src])
+    return output