D4CMPP2 0.4.0__py3-none-any.whl

D4CMPP2/networks/ISATPN_model.py

@@ -0,0 +1,199 @@
+import torch
+import torch.nn as nn
+
+from D4CMPP2.networks.base import (
+    Hyperparameter,
+    InputContract,
+    ISA_HYPERPARAMETERS,
+    ISA_OPTIMIZATION_SPACE,
+    MolecularNetwork,
+)
+from D4CMPP2.networks.src.ISAT import ISATconvolution_PM
+from D4CMPP2.networks.src.Linear import Linears
+from D4CMPP2.networks.src.BiDropout import Bi_Dropout
+from D4CMPP2.networks.src.GCN import graph_sum_pool
+from D4CMPP2.networks.src.pyg_hetero import relation_graph
+
+class ISATPN(MolecularNetwork):
+    """Positive-negative ISAT with contribution-distribution regularization."""
+
+    model_name = "isatpn"
+    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_feature", "get_score"),
+    )
+    hyperparameters = {
+        **ISA_HYPERPARAMETERS,
+        "score_variance_target": Hyperparameter(
+            "float", default=0.17, low=0.0, search_low=0.01,
+            search_high=1.0, grid=(0.1, 0.17, 0.25),
+            description="Target variance of positive and negative scores.",
+        ),
+        "score_mean_target": Hyperparameter(
+            "float", default=0.5, low=0.0, high=1.0,
+            search_low=0.0, search_high=1.0, grid=(0.3, 0.5, 0.7),
+            description="Target mean of positive and negative scores.",
+        ),
+        "score_magnitude_weight": Hyperparameter(
+            "float", default=0.09, low=0.0, search_low=0.0,
+            search_high=0.2, grid=(0.0, 0.03, 0.09, 0.15),
+            description="Weight of score magnitude regularization.",
+        ),
+        "score_variance_weight": Hyperparameter(
+            "float", default=0.03, low=0.0, search_low=0.0,
+            search_high=0.2, grid=(0.0, 0.03, 0.1),
+            description="Weight of score variance regularization.",
+        ),
+        "score_mean_weight": Hyperparameter(
+            "float", default=0.03, low=0.0, search_low=0.0,
+            search_high=0.2, grid=(0.0, 0.03, 0.1),
+            description="Weight of score mean regularization.",
+        ),
+        "bidirectional_dropout": Hyperparameter(
+            "float", default=0.4, low=0.0, high=0.5,
+            search_low=0.0, search_high=0.5,
+            grid=(0.0, 0.2, 0.4),
+            description="Probability of suppressing either contribution branch.",
+        ),
+    }
+    default_optimization_space = (
+        *ISA_OPTIMIZATION_SPACE,
+        "score_magnitude_weight",
+        "score_variance_weight",
+        "score_mean_weight",
+        "bidirectional_dropout",
+    )
+
+    def __init__(self, config):
+        aliases = {
+            "score_variance_target": "score_var",
+            "score_mean_target": "score_mean",
+            "score_magnitude_weight": "alpha",
+            "score_variance_weight": "beta",
+            "score_mean_weight": "gamma",
+        }
+        normalized = dict(config)
+        for new_name, old_name in aliases.items():
+            if new_name not in normalized and old_name in normalized:
+                normalized[new_name] = normalized[old_name]
+        super().__init__(normalized)
+        if self.config["target_dim"] != 1:
+            raise ValueError(
+                "ISATPN requires target_dim=1 because positive and negative "
+                "branches are subtracted into one prediction."
+            )
+        
+        
+        hidden_dim = self.config["hidden_dim"]
+        gcn_layers = self.config["conv_layers"]
+        dropout = self.config["dropout"]
+        linear_layers = self.config["linear_layers"]
+        last_linear_dim = int(hidden_dim/(2**linear_layers))
+        if last_linear_dim < 1:
+            raise ValueError(
+                "ISATPN requires hidden_dim >= 2 ** linear_layers; got "
+                f"hidden_dim={hidden_dim}, linear_layers={linear_layers}."
+            )
+
+        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.ISATconv_PM = ISATconvolution_PM(hidden_dim, hidden_dim, hidden_dim, nn.LeakyReLU(), gcn_layers, dropout, False, True, 0.1)
+
+        self.linears = Linears(hidden_dim,last_linear_dim, nn.ReLU(), linear_layers, dropout, False, False)
+        self.out_linear = nn.Linear(last_linear_dim, self.config["target_dim"])
+
+        self.Bi_Dropout=Bi_Dropout(self.config["bidirectional_dropout"])
+
+        self.score_var = self.config["score_variance_target"]
+        self.score_mean = self.config["score_mean_target"]
+        self.alpha = self.config["score_magnitude_weight"]
+        self.beta = self.config["score_variance_weight"]
+        self.gamma = self.config["score_mean_weight"]
+        
+    
+    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"ISATPN input is missing required fields {missing!r}.")
+
+        self.Bi_Dropout.set_drop_p(self.training)
+        
+        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)
+        
+        real_graph = relation_graph(graph, 'r_nd', 'r2r')
+
+        r_node_P, r_node_M, score_P, score_M = self.ISATconv_PM(graph, r_node, r_edge, i_node, d_edge, **kargs)
+        if kargs.get('get_feature',False):
+            return {'feature_P':r_node_P,'feature_N':r_node_M,'positive':score_P,'negative':score_M}
+
+        h_P = graph_sum_pool(real_graph, r_node_P)
+        h_P = self.linears(h_P)
+        h_P = self.out_linear(h_P)
+
+        h_M = graph_sum_pool(real_graph, r_node_M)
+        h_M = self.linears(h_M)
+        h_M = self.out_linear(h_M)
+
+        h=torch.concat([h_P,h_M],axis=1)
+        h=self.Bi_Dropout(h)
+        output=torch.unsqueeze(h[:,0]-h[:,1],axis=1)
+        if kargs.get('get_score',False):
+            return {'prediction':output, 'positive':score_P,'negative':score_M}
+        
+        self.p_score_var = self.ISATconv_PM.p_score_var
+        self.n_score_var = self.ISATconv_PM.n_score_var
+        self.p_score_mean = self.ISATconv_PM.p_score_mean
+        self.n_score_mean = self.ISATconv_PM.n_score_mean
+        self.p_score_ms = self.ISATconv_PM.p_score_ms
+        self.n_score_ms = self.ISATconv_PM.n_score_ms
+
+        return output
+        
+    def compute_loss(self, scores, targets):
+        mask = ~torch.isnan(targets)
+        if not torch.any(mask):
+            raise ValueError("A batch contains no finite target values for loss.")
+        scores,targets=self.Bi_Dropout.drop_label(scores[mask],targets[mask])
+        return torch.mean(torch.square(scores-targets))+\
+                (self.p_score_ms+self.n_score_ms)*self.alpha +\
+                (torch.abs(self.score_var-self.p_score_var)+torch.abs(self.score_var-self.n_score_var))*self.beta +\
+                (torch.abs(self.score_mean-self.p_score_mean)+torch.abs(self.score_mean-self.n_score_mean))*self.gamma
+
+
+network = ISATPN
+
+

D4CMPP2/networks/ISAT_model.py

@@ -0,0 +1,90 @@
+import torch.nn as nn
+
+from D4CMPP2.networks.base import (
+    InputContract,
+    ISA_HYPERPARAMETERS,
+    ISA_OPTIMIZATION_SPACE,
+    MolecularNetwork,
+)
+from D4CMPP2.networks.src.ISAT import ISATconvolution
+from D4CMPP2.networks.src.Linear import Linears
+from D4CMPP2.networks.src.GCN import graph_sum_pool
+from D4CMPP2.networks.src.pyg_hetero import relation_graph
+
+class ISAT(MolecularNetwork):
+    """Interpretable structure-attention network."""
+
+    model_name = "isat"
+    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 = ISA_HYPERPARAMETERS
+    default_optimization_space = ISA_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.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.ISATconv = ISATconvolution(hidden_dim, hidden_dim, hidden_dim, nn.LeakyReLU(), gcn_layers,dropout, False, True, 0.1)
+        
+        self.linears = Linears(hidden_dim, target_dim, nn.ReLU(), linear_layers, dropout, False, False, last=True)
+    
+    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"ISAT 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)
+        
+        r_node, score = self.ISATconv(graph, r_node, r_edge, i_node, d_edge)
+        h = graph_sum_pool(relation_graph(graph, 'r_nd', 'r2r'), r_node)
+        h = self.linears(h)
+
+        if kargs.get('get_score',False):
+            return {'prediction':h, 'positive':score}
+
+        return h
+
+network = ISAT

D4CMPP2/networks/MPNN_model.py

@@ -0,0 +1,56 @@
+import torch.nn as nn
+from D4CMPP2.networks.base import (
+    InputContract,
+    MolecularNetwork,
+    STANDARD_GRAPH_HYPERPARAMETERS,
+    STANDARD_GRAPH_OPTIMIZATION_SPACE,
+)
+from D4CMPP2.networks.src.MPNN import MPNNs
+from D4CMPP2.networks.src.GCN import graph_sum_pool
+from D4CMPP2.networks.src.Linear import Linears
+
+class MPNN(MolecularNetwork):
+    """Edge-aware message-passing network with graph-sum readout."""
+
+    model_name = "mpnn"
+    required_config = ("node_dim", "edge_dim", "target_dim")
+    input_contract = InputContract(
+        required=(
+            "compound_graphs",
+            "compound_node_feature",
+            "compound_edge_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"]
+        mpnn_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.edge_embedding = nn.Linear(self.config["edge_dim"], hidden_dim)
+
+        self.MPNNs = MPNNs(hidden_dim, hidden_dim, hidden_dim, hidden_dim, nn.LeakyReLU(), mpnn_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.LeakyReLU(), linear_layers, dropout, False, False, last=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"]
+        edge_feats = kwargs["compound_edge_feature"]
+
+        h = self.node_embedding(node_feats)
+        e = self.edge_embedding(edge_feats)
+        h = self.MPNNs(graph, h, e)
+        h = graph_sum_pool(graph, h)
+        h = self.Linears(h)
+        return h
+
+network = MPNN

D4CMPP2/networks/MPNNwithSolv_model.py

@@ -0,0 +1,72 @@
+import torch.nn as nn
+from D4CMPP2.networks.base import (
+    InputContract,
+    MolecularNetwork,
+    STANDARD_SOLVENT_HYPERPARAMETERS,
+    STANDARD_SOLVENT_OPTIMIZATION_SPACE,
+)
+from D4CMPP2.networks.src.MPNN import MPNNs
+from D4CMPP2.networks.src.GCN import graph_sum_pool
+from D4CMPP2.networks.src.Linear import Linears
+from D4CMPP2.networks.src.SolventLayer import SolventLayer
+
+class SolventMPNN(MolecularNetwork):
+    """Edge-aware compound MPNN coupled to a solvent graph branch."""
+
+    model_name = "mpnn_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"]
+        mpnn_layers = self.config["conv_layers"]
+        dropout = self.config["dropout"]
+        linear_layers = self.config["linear_layers"]
+
+        self.node_embedding = nn.Linear(self.config["node_dim"], hidden_dim)
+        self.edge_embedding = nn.Linear(self.config["edge_dim"], hidden_dim)
+        self.MPNNs = MPNNs(hidden_dim, hidden_dim, hidden_dim, hidden_dim, nn.ReLU(), mpnn_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,hidden_dim, nn.ReLU(), linear_layers, dropout, False, False) # in_feats, out_feats, activation, n_layers, dropout=0.2, batch_norm=False, residual_sum = False):
+        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'))
+        values = {
+            "compound_graphs": graph,
+            "compound_node_feature": node_feats,
+            "compound_edge_feature": edge_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"SolventMPNN input is missing required fields {missing!r}."
+            )
+
+        h = self.node_embedding(node_feats)
+        e = self.edge_embedding(edge_feats)
+        h = self.MPNNs(graph, h, e)
+        h = graph_sum_pool(graph, h)
+        h = self.Linears(h)
+        h = self.solvlayer(h, solv_graph, solv_node_feats)
+        return h
+
+network = SolventMPNN

D4CMPP2/networks/__init__.py

@@ -0,0 +1,25 @@
+"""Public contracts for built-in and custom D4CMPP2 networks."""
+
+from D4CMPP2.networks.base import (
+    Hyperparameter,
+    InputContract,
+    ModelConfig,
+    MolecularNetwork,
+)
+from D4CMPP2.networks.registry import (
+    ModelDefinition,
+    get_model,
+    register_network,
+    registered_models,
+)
+
+__all__ = [
+    "Hyperparameter",
+    "InputContract",
+    "ModelConfig",
+    "ModelDefinition",
+    "MolecularNetwork",
+    "get_model",
+    "register_network",
+    "registered_models",
+]

D4CMPP2/networks/base.py

@@ -0,0 +1,250 @@
+"""Small, explicit contracts shared by D4CMPP2 molecular networks."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import ClassVar, Mapping
+
+import torch
+import torch.nn as nn
+
+
+_MISSING = object()
+
+
+@dataclass(frozen=True)
+class Hyperparameter:
+    """One validated model field and its optional optimization domain."""
+
+    kind: str
+    default: object = _MISSING
+    low: float | int | None = None
+    high: float | int | None = None
+    search_low: float | int | None = None
+    search_high: float | int | None = None
+    step: float | int | None = None
+    values: tuple[object, ...] = ()
+    scale: str = "linear"
+    grid: tuple[object, ...] = ()
+    description: str = ""
+
+    def validate(self, name, value):
+        if self.kind == "int":
+            if isinstance(value, bool) or not isinstance(value, int):
+                raise TypeError(f"{name} must be an integer, got {value!r}.")
+        elif self.kind == "float":
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
+                raise TypeError(f"{name} must be a number, got {value!r}.")
+            value = float(value)
+        elif self.kind == "categorical":
+            if value not in self.values:
+                raise ValueError(
+                    f"{name} must be one of {list(self.values)!r}, got {value!r}."
+                )
+            return value
+        else:
+            raise ValueError(
+                f"Hyperparameter {name!r} has unknown kind {self.kind!r}."
+            )
+        if self.low is not None and value < self.low:
+            raise ValueError(f"{name} must be >= {self.low}, got {value!r}.")
+        if self.high is not None and value > self.high:
+            raise ValueError(f"{name} must be <= {self.high}, got {value!r}.")
+        return value
+
+
+STANDARD_GRAPH_HYPERPARAMETERS = MappingProxyType(
+    {
+        "hidden_dim": Hyperparameter(
+            "int", default=64, low=1, search_low=16, search_high=256,
+            step=16, grid=(32, 64, 128, 256),
+            description="Hidden 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 message-passing 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.",
+        ),
+    }
+)
+STANDARD_GRAPH_OPTIMIZATION_SPACE = tuple(STANDARD_GRAPH_HYPERPARAMETERS)
+STANDARD_SOLVENT_HYPERPARAMETERS = MappingProxyType(
+    {
+        **STANDARD_GRAPH_HYPERPARAMETERS,
+        "solvent_hidden_dim": Hyperparameter(
+            "int", default=64, low=1, search_low=16, search_high=256,
+            step=16, grid=(32, 64, 128),
+            description="Hidden width of the solvent graph branch.",
+        ),
+        "solvent_conv_layers": Hyperparameter(
+            "int", default=4, low=1, search_low=1, search_high=8,
+            step=1, grid=(2, 4, 6),
+            description="Number of solvent graph convolution layers.",
+        ),
+        "solvent_linear_layers": Hyperparameter(
+            "int", default=2, low=1, search_low=1, search_high=5,
+            step=1, grid=(1, 2, 3),
+            description="Number of solvent projection layers.",
+        ),
+        "solvent_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 in the solvent branch.",
+        ),
+    }
+)
+STANDARD_SOLVENT_OPTIMIZATION_SPACE = tuple(STANDARD_SOLVENT_HYPERPARAMETERS)
+STANDARD_SOLVENT_CONFIG_ALIASES = MappingProxyType(
+    {
+        "solv_hidden_dim": "solvent_hidden_dim",
+        "solv_conv_layers": "solvent_conv_layers",
+        "solv_linear_layers": "solvent_linear_layers",
+        "solv_dropout": "solvent_dropout",
+    }
+)
+ISA_HYPERPARAMETERS = MappingProxyType(
+    {
+        **STANDARD_GRAPH_HYPERPARAMETERS,
+        "linear_layers": Hyperparameter(
+            "int", default=4, low=1, high=4,
+            search_low=1, search_high=4, step=1, grid=(1, 2, 3, 4),
+            description="Number of prediction-head linear layers.",
+        ),
+        "dropout": Hyperparameter(
+            "float", default=0.1, 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.",
+        ),
+    }
+)
+ISA_OPTIMIZATION_SPACE = tuple(ISA_HYPERPARAMETERS)
+
+
+class ModelConfig(Mapping):
+    """Immutable, model-only view of the larger training configuration."""
+
+    def __init__(self, values):
+        self._values = MappingProxyType(dict(values))
+
+    def __getitem__(self, key):
+        return self._values[key]
+
+    def __iter__(self):
+        return iter(self._values)
+
+    def __len__(self):
+        return len(self._values)
+
+    def __repr__(self):
+        return f"ModelConfig({dict(self._values)!r})"
+
+
+@dataclass(frozen=True)
+class InputContract:
+    required: tuple[str, ...]
+    optional: tuple[str, ...] = ()
+
+    def validate(self, batch):
+        missing = [name for name in self.required if name not in batch]
+        if missing:
+            raise ValueError(
+                f"Network input is missing required fields {missing!r}. "
+                f"Required fields: {list(self.required)!r}."
+            )
+
+
+def masked_mse(prediction, target):
+    mask = ~torch.isnan(target)
+    if not torch.any(mask):
+        raise ValueError("A batch contains no finite target values for loss.")
+    return torch.mean((prediction[mask] - target[mask]) ** 2)
+
+
+class MolecularNetwork(nn.Module, ABC):
+    """Base contract for trainable molecular property networks."""
+
+    model_name: ClassVar[str]
+    input_contract: ClassVar[InputContract]
+    required_config: ClassVar[tuple[str, ...]] = (
+        "node_dim",
+        "target_dim",
+    )
+    hyperparameters: ClassVar[Mapping[str, Hyperparameter]] = {}
+    default_optimization_space: ClassVar[tuple[str, ...]] = ()
+    config_aliases: ClassVar[Mapping[str, str]] = STANDARD_SOLVENT_CONFIG_ALIASES
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = self.validate_config(config)
+
+    @classmethod
+    def validate_config(cls, config):
+        normalized = dict(config)
+        for old_name, canonical_name in cls.config_aliases.items():
+            if old_name in normalized:
+                if (
+                    canonical_name in normalized
+                    and normalized[canonical_name] != normalized[old_name]
+                ):
+                    raise ValueError(
+                        f"{cls.__name__} received conflicting values for legacy "
+                        f"{old_name!r} and canonical {canonical_name!r}."
+                    )
+                normalized.setdefault(canonical_name, normalized[old_name])
+        missing = [name for name in cls.required_config if name not in normalized]
+        if missing:
+            raise ValueError(
+                f"{cls.__name__} config is missing required fields {missing!r}."
+            )
+        values = {name: normalized[name] for name in cls.required_config}
+        for name, field in cls.hyperparameters.items():
+            if name in normalized and normalized[name] is not None:
+                value = normalized[name]
+            elif field.default is not _MISSING:
+                value = field.default
+            else:
+                raise ValueError(
+                    f"{cls.__name__} config requires hyperparameter {name!r}."
+                )
+            values[name] = field.validate(name, value)
+        return ModelConfig(values)
+
+    @classmethod
+    def optimization_space(cls, keys=None):
+        selected = (
+            cls.default_optimization_space if keys is None else tuple(keys)
+        )
+        unknown = [name for name in selected if name not in cls.hyperparameters]
+        if unknown:
+            raise ValueError(
+                f"{cls.model_name} does not define hyperparameters {unknown!r}. "
+                f"Available keys: {sorted(cls.hyperparameters)!r}."
+            )
+        return {name: cls.hyperparameters[name] for name in selected}
+
+    def validate_input(self, batch):
+        self.input_contract.validate(batch)
+
+    def compute_loss(self, prediction, target):
+        return masked_mse(prediction, target)
+
+    def loss_fn(self, prediction, target):
+        """Temporary training-manager bridge during the staged migration."""
+        return self.compute_loss(prediction, target)
+
+    @abstractmethod
+    def forward(self, **batch):
+        raise NotImplementedError