D4CMPP2 0.4.0__py3-none-any.whl

D4CMPP2/networks/DMPNNwithSolv_model.py

@@ -0,0 +1,89 @@
+
+import torch.nn as nn
+from D4CMPP2.networks.base import (
+    InputContract,
+    MolecularNetwork,
+    STANDARD_SOLVENT_HYPERPARAMETERS,
+    STANDARD_SOLVENT_OPTIMIZATION_SPACE,
+)
+from D4CMPP2.networks.src.GCN import graph_sum_pool
+
+from D4CMPP2.networks.src.SolventLayer import SolventLayer
+from D4CMPP2.networks.src.DMPNN import DMPNNLayer
+
+
+class SolventDMPNN(MolecularNetwork):
+    """Directed-edge compound encoder coupled to a solvent graph branch."""
+
+    model_name = "dmpnn_solvent"
+    required_config = ("node_dim", "edge_dim", "target_dim")
+    input_contract = InputContract(
+        required=(
+            "compound_graphs",
+            "compound_node_feature",
+            "compound_edge_feature",
+            "solvent_graphs",
+            "solvent_node_feature",
+        )
+    )
+    hyperparameters = STANDARD_SOLVENT_HYPERPARAMETERS
+    default_optimization_space = STANDARD_SOLVENT_OPTIMIZATION_SPACE
+
+    def __init__(self, config):
+        super().__init__(config)
+        hidden_dim = self.config["hidden_dim"]
+        conv_layers = self.config["conv_layers"]
+        dropout = self.config["dropout"]
+
+        self.embedding_node_lin = nn.Linear(self.config["node_dim"], hidden_dim, bias=True)
+        self.embedding_edge_lin = nn.Linear(self.config["edge_dim"], hidden_dim, bias=True)
+        self.init_h_func = nn.Sequential(
+            nn.Linear(2 * hidden_dim, hidden_dim, bias=True),
+            nn.LeakyReLU()
+        )
+        self.W_a = nn.Sequential(
+            nn.Linear(2 * hidden_dim, hidden_dim, bias=True),
+            nn.LeakyReLU()
+        )
+        self.dropout_layer = nn.Dropout(dropout)
+        self.DMPNNLayer = nn.ModuleList([DMPNNLayer(hidden_dim,hidden_dim,hidden_dim,nn.LeakyReLU(),dropout) for _ in range(conv_layers)])
+        self.SolvLayer = SolventLayer(self.config)
+        
+
+    def forward(self, **kwargs):
+        graph = kwargs.get('compound_graphs', kwargs.get('graph'))
+        node_feats = kwargs.get('compound_node_feature', kwargs.get('node_feats'))
+        edge_feats = kwargs.get('compound_edge_feature', kwargs.get('edge_feats'))
+        solv_graph = kwargs.get('solvent_graphs', kwargs.get('solv_graph'))
+        solv_node_feats = kwargs.get('solvent_node_feature', kwargs.get('solv_node_feats'))
+        missing = [
+            name
+            for name, value in {
+                "compound_graphs": graph,
+                "compound_node_feature": node_feats,
+                "compound_edge_feature": edge_feats,
+                "solvent_graphs": solv_graph,
+                "solvent_node_feature": solv_node_feats,
+            }.items()
+            if value is None
+        ]
+        if missing:
+            raise ValueError(
+                f"SolventDMPNN input is missing required fields {missing!r}."
+            )
+
+        node = self.embedding_node_lin(node_feats)
+        edge = self.embedding_edge_lin(edge_feats)
+
+        direct_feats = None
+        backward_feats = None
+        for layer in self.DMPNNLayer:
+            hidden_feats, direct_feats, backward_feats = layer(graph, node, edge, direct_feats, backward_feats)
+
+        graph_feats = graph_sum_pool(graph, hidden_feats)
+        output = self.SolvLayer(graph_feats, solv_graph, solv_node_feats)
+        return output
+
+network = SolventDMPNN
+
+

D4CMPP2/networks/GAT_model.py

@@ -0,0 +1,48 @@
+import torch.nn as nn
+from D4CMPP2.networks.base import (
+    InputContract,
+    MolecularNetwork,
+    STANDARD_GRAPH_HYPERPARAMETERS,
+    STANDARD_GRAPH_OPTIMIZATION_SPACE,
+)
+from D4CMPP2.networks.src.GAT import GATs
+from D4CMPP2.networks.src.GCN import graph_sum_pool
+from D4CMPP2.networks.src.Linear import Linears
+
+class GAT(MolecularNetwork):
+    """Graph-attention network with graph-sum readout."""
+
+    model_name = "gat"
+    input_contract = InputContract(
+        required=("compound_graphs", "compound_node_feature")
+    )
+    hyperparameters = STANDARD_GRAPH_HYPERPARAMETERS
+    default_optimization_space = STANDARD_GRAPH_OPTIMIZATION_SPACE
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        hidden_dim = self.config["hidden_dim"]
+        gcn_layers = self.config["conv_layers"]
+        dropout = self.config["dropout"]
+        linear_layers = self.config["linear_layers"]
+        target_dim = self.config["target_dim"]
+
+
+        self.node_embedding = nn.Linear(self.config["node_dim"], hidden_dim)
+
+        self.GATs = GATs(hidden_dim, hidden_dim, hidden_dim, nn.ReLU(), gcn_layers, dropout, False, True) # in_feats, hidden_feats, out_feats, activation, n_layers, dropout=0.2, batch_norm=False, residual_sum = False):
+        self.Linears = Linears(hidden_dim,target_dim, nn.ReLU(), linear_layers, dropout, False, False, True) # in_feats, out_feats, activation, n_layers, dropout=0.2, batch_norm=False, residual_sum = False):
+
+    def forward(self, **kwargs):
+        self.validate_input(kwargs)
+        graph = kwargs["compound_graphs"]
+        node_feats = kwargs["compound_node_feature"]
+
+        h = self.node_embedding(node_feats)
+        h = self.GATs(graph, h)
+        h = graph_sum_pool(graph, h)
+        h = self.Linears(h)
+        return h
+
+network = GAT

D4CMPP2/networks/GATwithSolv_model.py

@@ -0,0 +1,63 @@
+import torch.nn as nn
+from D4CMPP2.networks.base import (
+    InputContract,
+    MolecularNetwork,
+    STANDARD_SOLVENT_HYPERPARAMETERS,
+    STANDARD_SOLVENT_OPTIMIZATION_SPACE,
+)
+from D4CMPP2.networks.src.GAT import GATs
+from D4CMPP2.networks.src.GCN import graph_sum_pool
+from D4CMPP2.networks.src.SolventLayer import SolventLayer
+
+class SolventGAT(MolecularNetwork):
+    """Graph-attention compound encoder coupled to a solvent branch."""
+
+    model_name = "gat_solvent"
+    input_contract = InputContract(
+        required=(
+            "compound_graphs",
+            "compound_node_feature",
+            "solvent_graphs",
+            "solvent_node_feature",
+        )
+    )
+    hyperparameters = STANDARD_SOLVENT_HYPERPARAMETERS
+    default_optimization_space = STANDARD_SOLVENT_OPTIMIZATION_SPACE
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        hidden_dim = self.config["hidden_dim"]
+        gcn_layers = self.config["conv_layers"]
+        dropout = self.config["dropout"]
+
+
+        self.node_embedding = nn.Linear(self.config["node_dim"], hidden_dim)
+
+        self.GATs = GATs(hidden_dim, hidden_dim, hidden_dim, nn.ReLU(), gcn_layers, dropout, False, True)
+        self.SolventLayer = SolventLayer(self.config)
+
+    def forward(self, **kwargs):
+        graph = kwargs.get('compound_graphs', kwargs.get('graph'))
+        node_feats = kwargs.get('compound_node_feature', kwargs.get('node_feats'))
+        solv_graph = kwargs.get('solvent_graphs', kwargs.get('solv_graph'))
+        solv_node_feats = kwargs.get('solvent_node_feature', kwargs.get('solv_node_feats'))
+        values = {
+            "compound_graphs": graph,
+            "compound_node_feature": node_feats,
+            "solvent_graphs": solv_graph,
+            "solvent_node_feature": solv_node_feats,
+        }
+        missing = [name for name, value in values.items() if value is None]
+        if missing:
+            raise ValueError(
+                f"SolventGAT input is missing required fields {missing!r}."
+            )
+
+        h = self.node_embedding(node_feats)
+        h = self.GATs(graph, h)
+        h = graph_sum_pool(graph, h)
+        h = self.SolventLayer(h,solv_graph,solv_node_feats)
+        return h
+
+network = SolventGAT

D4CMPP2/networks/GCN_model.py

@@ -0,0 +1,113 @@
+import torch.nn as nn
+
+from D4CMPP2.networks.base import (
+    Hyperparameter,
+    InputContract,
+    MolecularNetwork,
+)
+from D4CMPP2.networks.src.GCN import GCNs, graph_sum_pool
+from D4CMPP2.networks.src.Linear import Linears
+
+
+class GCN(MolecularNetwork):
+    """Graph convolution using transformed self-plus-neighbor sums."""
+
+    model_name = "gcn"
+    input_contract = InputContract(
+        required=("compound_graphs", "compound_node_feature"),
+    )
+    hyperparameters = {
+        "hidden_dim": Hyperparameter(
+            "int",
+            default=64,
+            low=1,
+            search_low=16,
+            search_high=256,
+            step=16,
+            grid=(32, 64, 128, 256),
+            description="Hidden node and graph representation width.",
+        ),
+        "conv_layers": Hyperparameter(
+            "int",
+            default=4,
+            low=1,
+            search_low=1,
+            search_high=8,
+            step=1,
+            grid=(2, 4, 6, 8),
+            description="Number of graph convolution layers.",
+        ),
+        "linear_layers": Hyperparameter(
+            "int",
+            default=2,
+            low=1,
+            search_low=1,
+            search_high=5,
+            step=1,
+            grid=(1, 2, 3, 4),
+            description="Number of prediction-head linear layers.",
+        ),
+        "dropout": Hyperparameter(
+            "float",
+            default=0.2,
+            low=0.0,
+            high=0.5,
+            search_low=0.0,
+            search_high=0.5,
+            grid=(0.0, 0.1, 0.2, 0.3, 0.5),
+            description="Dropout probability.",
+        ),
+    }
+    default_optimization_space = (
+        "hidden_dim",
+        "conv_layers",
+        "linear_layers",
+        "dropout",
+    )
+
+    def __init__(self, config):
+        super().__init__(config)
+        hidden_dim = self.config["hidden_dim"]
+        self.node_embedding = nn.Linear(self.config["node_dim"], hidden_dim)
+        # Keep the established module registration names so saved state dictionaries
+        # remain readable while the public class/config contract is redesigned.
+        self.GCNs = GCNs(
+            hidden_dim,
+            hidden_dim,
+            hidden_dim,
+            nn.ReLU(),
+            self.config["conv_layers"],
+            self.config["dropout"],
+            False,
+            True,
+        )
+        self.Linears = Linears(
+            hidden_dim,
+            self.config["target_dim"],
+            nn.ReLU(),
+            self.config["linear_layers"],
+            self.config["dropout"],
+            False,
+            False,
+            True,
+        )
+
+    @property
+    def convolutions(self):
+        return self.GCNs
+
+    @property
+    def prediction_head(self):
+        return self.Linears
+
+    def forward(self, **batch):
+        self.validate_input(batch)
+        graph = batch["compound_graphs"]
+        node = self.node_embedding(batch["compound_node_feature"])
+        node = self.GCNs(graph, node)
+        graph_feature = graph_sum_pool(graph, node)
+        return self.Linears(graph_feature)
+
+
+# Staged bridge for code paths not yet migrated to the typed registry.
+network = GCN

D4CMPP2/networks/GCNwithSolv_model.py

@@ -0,0 +1,103 @@
+import torch.nn as nn
+import torch
+from D4CMPP2.networks.base import (
+    Hyperparameter,
+    InputContract,
+    MolecularNetwork,
+    STANDARD_SOLVENT_HYPERPARAMETERS,
+    STANDARD_SOLVENT_OPTIMIZATION_SPACE,
+)
+from D4CMPP2.networks.src.GCN import GCNs, graph_sum_pool
+from D4CMPP2.networks.src.Linear import Linears
+
+class SolventGCN(MolecularNetwork):
+    """Two-branch GCN for compound-solvent property prediction."""
+
+    model_name = "gcn_solvent"
+    input_contract = InputContract(
+        required=(
+            "compound_graphs",
+            "compound_node_feature",
+            "solvent_graphs",
+            "solvent_node_feature",
+        )
+    )
+    hyperparameters = {
+        **STANDARD_SOLVENT_HYPERPARAMETERS,
+        "fusion_dim": Hyperparameter(
+            "int", default=16, low=1, search_low=8, search_high=128,
+            step=8, grid=(16, 32, 64),
+            description="Width of the fused compound-solvent representation.",
+        ),
+    }
+    default_optimization_space = (
+        *STANDARD_SOLVENT_OPTIMIZATION_SPACE,
+        "fusion_dim",
+    )
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        hidden_dim = self.config["hidden_dim"]
+        conv_layers = self.config["conv_layers"]
+        linear_layers = self.config["linear_layers"]
+        dropout = self.config["dropout"]
+        solvent_hidden_dim = self.config["solvent_hidden_dim"]
+        solvent_conv_layers = self.config["solvent_conv_layers"]
+        solvent_linear_layers = self.config["solvent_linear_layers"]
+        solvent_dropout = self.config["solvent_dropout"]
+        fusion_dim = self.config["fusion_dim"]
+        self.node_embedding = nn.Linear(self.config["node_dim"], hidden_dim)
+        self.node_embedding_solv = nn.Linear(
+            self.config["node_dim"], solvent_hidden_dim
+        )
+
+        self.GCNs = GCNs(hidden_dim, hidden_dim, hidden_dim, nn.ReLU(), conv_layers, dropout, False, True)
+        self.GCNs_solv = GCNs(solvent_hidden_dim, solvent_hidden_dim, solvent_hidden_dim, nn.ReLU(), solvent_conv_layers, solvent_dropout, False, True)
+        self.Linears1 = Linears(hidden_dim,hidden_dim, nn.ReLU(), 2, dropout, False, False)
+        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,fusion_dim, nn.ReLU(), linear_layers, dropout, False, False)
+
+        self.output = nn.Linear(fusion_dim, self.config["target_dim"])
+
+    def forward(self, **kwargs):
+        batch = {
+            "compound_graphs": kwargs.get(
+                "compound_graphs", kwargs.get("graph")
+            ),
+            "compound_node_feature": kwargs.get(
+                "compound_node_feature", kwargs.get("node_feats")
+            ),
+            "solvent_graphs": kwargs.get(
+                "solvent_graphs", kwargs.get("solv_graph")
+            ),
+            "solvent_node_feature": kwargs.get(
+                "solvent_node_feature", kwargs.get("solv_node_feats")
+            ),
+        }
+        missing = [name for name, value in batch.items() if value is None]
+        if missing:
+            raise ValueError(
+                f"SolventGCN input is missing required fields {missing!r}."
+            )
+        graph = batch["compound_graphs"]
+        node_feats = batch["compound_node_feature"]
+        solv_graph = batch["solvent_graphs"]
+        solv_node_feats = batch["solvent_node_feature"]
+
+        h = self.node_embedding(node_feats)
+        h = self.GCNs(graph, h)
+        h = graph_sum_pool(graph, h)
+
+        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 = self.Linears1(h)
+        h_solv = self.Linears2(h_solv)
+        h = torch.cat([h,h_solv],axis=1)
+        h = self.Linears3(h)
+        h = self.output(h)
+        return h
+
+network = SolventGCN

D4CMPP2/networks/GC_model.py

@@ -0,0 +1,122 @@
+import torch.nn as nn
+
+from D4CMPP2.networks.base import (
+    Hyperparameter,
+    InputContract,
+    MolecularNetwork,
+    STANDARD_GRAPH_HYPERPARAMETERS,
+    STANDARD_GRAPH_OPTIMIZATION_SPACE,
+)
+from D4CMPP2.networks.src.GC import GCconvolution
+from D4CMPP2.networks.src.GCN import graph_sum_pool
+from D4CMPP2.networks.src.pyg_hetero import relation_graph
+
+class GroupContributionNetwork(MolecularNetwork):
+    """Group-contribution network over the ISA heterogeneous graph."""
+
+    model_name = "group_contribution"
+    required_config = ("node_dim", "edge_dim", "target_dim")
+    input_contract = InputContract(
+        required=(
+            "compound_graphs",
+            "compound_r_node",
+            "compound_i_node",
+            "compound_r2r_edge",
+            "compound_d2d_edge",
+        ),
+        optional=("get_score",),
+    )
+    hyperparameters = {
+        **{
+            name: field
+            for name, field in STANDARD_GRAPH_HYPERPARAMETERS.items()
+            if name != "linear_layers"
+        },
+        "activation": Hyperparameter(
+            "categorical",
+            default="Tanh",
+            values=("Tanh", "Sigmoid", "ReLU", "Identity"),
+            grid=("Tanh", "Sigmoid", "ReLU"),
+            description="Activation applied to the normalized contribution sum.",
+        ),
+    }
+    default_optimization_space = (
+        *(
+            name
+            for name in STANDARD_GRAPH_OPTIMIZATION_SPACE
+            if name != "linear_layers"
+        ),
+        "activation",
+    )
+
+    def __init__(self, config):
+        if "activation" not in config and "func_f" in config:
+            config = {**config, "activation": config["func_f"]}
+        super().__init__(config)
+        if self.config["target_dim"] != 1:
+            raise ValueError(
+                "GroupContributionNetwork requires target_dim=1 because its "
+                "contribution head produces one scalar per molecule."
+            )
+
+        hidden_dim = self.config["hidden_dim"]
+        gcn_layers = self.config["conv_layers"]
+        dropout = self.config["dropout"]
+
+        self.embedding_rnode_lin = nn.Sequential(
+            nn.Linear(self.config["node_dim"], hidden_dim, bias=False)
+        )
+        self.embedding_inode_lin = nn.Sequential(
+            nn.Linear(1, hidden_dim, bias=False)
+        )
+        self.embedding_edge_lin = nn.Sequential(
+            nn.Linear(self.config["edge_dim"], hidden_dim, bias=False)
+        )
+
+        self.GCconv = GCconvolution(hidden_dim, hidden_dim, hidden_dim, nn.LeakyReLU(), gcn_layers, dropout, False, True, 0.1)
+        
+        self.linear = nn.Linear(1,1)
+        self.normalize = nn.LayerNorm(1)
+        self.func_f = getattr(nn, self.config["activation"])()
+    
+    def forward(self, **kargs):
+        graph = kargs.get('compound_graphs', kargs.get('graph'))
+        r_node = kargs.get('compound_r_node', kargs.get('r_node'))
+        i_node = kargs.get('compound_i_node', kargs.get('i_node'))
+        r_edge = kargs.get('compound_r2r_edge', kargs.get('r_edge'))
+        d_edge = kargs.get('compound_d2d_edge', kargs.get('d_edge'))
+        missing = [
+            name
+            for name, value in {
+                "compound_graphs": graph,
+                "compound_r_node": r_node,
+                "compound_i_node": i_node,
+                "compound_r2r_edge": r_edge,
+                "compound_d2d_edge": d_edge,
+            }.items()
+            if value is None
+        ]
+        if missing:
+            raise ValueError(
+                f"GroupContributionNetwork input is missing required fields {missing!r}."
+            )
+
+        r_node = r_node.float()
+        r_node = self.embedding_rnode_lin(r_node)
+        i_node = i_node.float()
+        i_node = self.embedding_inode_lin(i_node)
+        r_edge = r_edge.float()
+        r_edge = self.embedding_edge_lin(r_edge)
+        
+        score = self.GCconv(graph, r_node, r_edge, i_node, d_edge)
+        h = graph_sum_pool(relation_graph(graph, 'd_nd', 'd2d'), score)
+        h = self.normalize(h)
+        h = self.func_f(h)
+        h = self.linear(h)
+        
+        if kargs.get('get_score',False):
+            return {'prediction':h, 'positive':score}
+
+        return h
+
+network = GroupContributionNetwork

D4CMPP2/networks/ISATPM_model.py

@@ -0,0 +1,14 @@
+"""Deprecated compatibility wrapper for the former ISATPM module name."""
+
+import warnings
+
+warnings.warn(
+    "D4CMPP2.networks.ISATPM_model is deprecated; use "
+    "D4CMPP2.networks.ISATPN_model instead.",
+    FutureWarning,
+    stacklevel=2,
+)
+
+from D4CMPP2.networks.ISATPN_model import ISATPN, network
+
+__all__ = ["ISATPN", "network"]