k3-node 1.0.0__py3-none-any.whl

k3_node/models/__init__.py

@@ -0,0 +1,424 @@
+r"""k3-node ports of `torch_geometric.nn.models`."""
+
+from .mlp import MLP
+from .attract_repel import ARLinkPredictor
+from .autoencoder import InnerProductDecoder, GAE, VGAE, ARGA, ARGVA
+from .deep_graph_infomax import DeepGraphInfomax
+from .deepgcn import DeepGCNLayer
+from .attentive_fp import AttentiveFP
+from .jumping_knowledge import JumpingKnowledge, HeteroJumpingKnowledge
+from .mask_label import MaskLabel
+from .meta import MetaLayer
+from .pmlp import PMLP
+from .polynormer import Polynormer
+from .basic_gnn import BasicGNN, GCN, GraphSAGE, GIN, GAT, PNA, EdgeCNN
+from .label_prop import LabelPropagation
+from .correct_and_smooth import CorrectAndSmooth
+from .lightgcn import LightGCN, BPRLoss
+from .linkx import LINKX, SparseLinear
+from .rect import RECT_L
+from .signed_gcn import SignedGCN
+from .neural_fingerprint import NeuralFingerprint
+from .graph_unet import GraphUNet
+from .rev_gnn import GroupAddRev
+from .sgformer import SGFormer
+from .node2vec import Node2Vec
+from .metapath2vec import MetaPath2Vec
+from .renet import RENet
+from .tgn import (
+    TGNMemory,
+    IdentityMessage,
+    LastAggregator,
+    MeanAggregator,
+    TimeEncoder,
+    LastNeighborLoader,
+)
+from .schnet import (
+    SchNet,
+    CFConv,
+    InteractionBlock as SchNetInteractionBlock,
+    GaussianSmearing,
+    ShiftedSoftplus,
+    RadiusInteractionGraph,
+)
+from .dimenet import (
+    DimeNet,
+    DimeNetPlusPlus,
+    BesselBasisLayer,
+    SphericalBasisLayer,
+    triplets,
+)
+from .gnnff import (
+    GNNFF,
+    NodeBlock,
+    EdgeBlock,
+    GaussianFilter,
+)
+from .gpse import (
+    GPSE,
+    GPSENodeEncoder,
+    GeneralLayer,
+    GeneralMultiLayer,
+    GNNStackStage,
+    GNNInductiveHybridMultiHead,
+)
+from .visnet import ViSNet
+from .lpformer import LPFormer, LPAttLayer
+from .graphmae2 import (
+    GraphMAE2,
+    sce_loss,
+    load_graphmae2_weights,
+    download_graphmae2_checkpoint,
+)
+from .graphormer import (
+    Graphormer,
+    GraphNodeFeature,
+    GraphAttnBias,
+    GraphormerMultiheadAttention,
+    GraphormerGraphEncoderLayer,
+    GraphormerGraphEncoder,
+    load_graphormer_weights,
+    download_graphormer_checkpoint,
+)
+from .graphormer_3d import (
+    Graphormer3D,
+    GaussianLayer,
+    RBF,
+    Graphormer3DEncoderLayer,
+    NodeTaskHead,
+    load_graphormer3d_weights,
+    download_graphormer3d_checkpoint,
+)
+from .captum import to_captum_model, to_captum_input, captum_output_to_dicts
+from .gps_model import (
+    GPSModel,
+    GPSLayer,
+    CustomGatedGCN,
+    AtomEncoder,
+    BondEncoder,
+    RWSEEncoder,
+    SANGraphHead,
+    load_gps_weights,
+    download_gps_checkpoint,
+)
+from .grover import (
+    GROVER,
+    GTransEncoder,
+    Readout,
+    load_grover_weights,
+    download_grover_checkpoint,
+)
+from .mole_bert import (
+    MoleBERT,
+    MoleBERTGNN,
+    MoleBERTGINConv,
+    load_mole_bert_weights,
+    download_mole_bert_checkpoint,
+)
+from .unimol import (
+    UniMolModel,
+    UniMolConfGenModel,
+    UniMolDockingModel,
+    GaussianLayer as UniMolGaussianLayer,
+    NumericalEmbed as UniMolNumericalEmbed,
+    NonLinearHead as UniMolNonLinearHead,
+    DistanceHead as UniMolDistanceHead,
+    ClassificationHead as UniMolClassificationHead,
+    LinearHead as UniMolLinearHead,
+    MaskLMHead as UniMolMaskLMHead,
+    download_unimol_checkpoint,
+    load_unimol_weights,
+)
+from .unimol2 import (
+    UniMol2Model,
+    AtomFeature as UniMol2AtomFeature,
+    EdgeFeature as UniMol2EdgeFeature,
+    SE3InvariantKernel as UniMol2SE3Kernel,
+    MovementPredictionHead as UniMol2MovementHead,
+    download_unimol2_checkpoint,
+    load_unimol2_weights,
+)
+from .unimol_plus import (
+    UniMolPlusPCQModel,
+    UniMolPlusOC20Model,
+    EnergyHead as UniMolPlusEnergyHead,
+    download_unimol_plus_checkpoint,
+    load_unimol_plus_weights,
+)
+from .unimol_docking_v2 import (
+    DockingPoseModelV2,
+    download_unimol_docking_checkpoint,
+    load_unimol_docking_weights,
+)
+from . import materials
+from . import bio
+from . import chemistry
+from .materials import (
+    MEGNet,
+    MEGNetBlock,
+    MEGNetGraphConv,
+    M3GNet,
+    M3GNetBlock,
+    M3GNetGraphConv,
+    ThreeBodyInteractions,
+    TensorNet,
+    TensorEmbedding,
+    TensorNetInteraction,
+    CHGNet,
+    CHGNetAtomGraphBlock,
+    CHGNetBondGraphBlock,
+    SO3Net,
+    SO3Convolution,
+    RealSphericalHarmonics,
+    GRACE,
+    GraceSPBasis,
+    GraceACEStack,
+    QET,
+    LinearQeq,
+    ElectrostaticPotential,
+    TransformedTargetModel,
+    Potential,
+    BondExpansion as MatGLBondExpansion,
+    GaussianExpansion as MatGLGaussianExpansion,
+    RadialBesselFunction as MatGLRadialBesselFunction,
+    FourierExpansion as MatGLFourierExpansion,
+    ChebyshevRadialBasis as MatGLChebyshevRadialBasis,
+    SphericalBesselFunction as MatGLSphericalBesselFunction,
+    SphericalBesselWithHarmonics as MatGLSphericalBesselWithHarmonics,
+    ReduceReadOut as MatGLReduceReadOut,
+    WeightedReadOut as MatGLWeightedReadOut,
+    WeightedAtomReadOut as MatGLWeightedAtomReadOut,
+    Set2SetReadOut as MatGLSet2SetReadOut,
+    EdgeSet2Set as MatGLEdgeSet2Set,
+    download_matgl_checkpoint,
+    load_matgl_weights,
+    load_model as load_matgl_model,
+    get_available_pretrained_models as get_available_matgl_models,
+)
+
+__all__ = [
+    "materials",
+    "bio",
+    "chemistry",
+    "MEGNet",
+    "MEGNetBlock",
+    "MEGNetGraphConv",
+    "M3GNet",
+    "M3GNetBlock",
+    "M3GNetGraphConv",
+    "ThreeBodyInteractions",
+    "TensorNet",
+    "TensorEmbedding",
+    "TensorNetInteraction",
+    "CHGNet",
+    "CHGNetAtomGraphBlock",
+    "CHGNetBondGraphBlock",
+    "SO3Net",
+    "SO3Convolution",
+    "RealSphericalHarmonics",
+    "GRACE",
+    "GraceSPBasis",
+    "GraceACEStack",
+    "QET",
+    "LinearQeq",
+    "ElectrostaticPotential",
+    "TransformedTargetModel",
+    "Potential",
+    "MatGLBondExpansion",
+    "MatGLGaussianExpansion",
+    "MatGLRadialBesselFunction",
+    "MatGLFourierExpansion",
+    "MatGLChebyshevRadialBasis",
+    "MatGLSphericalBesselFunction",
+    "MatGLSphericalBesselWithHarmonics",
+    "MatGLReduceReadOut",
+    "MatGLWeightedReadOut",
+    "MatGLWeightedAtomReadOut",
+    "MatGLSet2SetReadOut",
+    "MatGLEdgeSet2Set",
+    "download_matgl_checkpoint",
+    "load_matgl_weights",
+    "load_matgl_model",
+    "get_available_matgl_models",
+    "MLP",
+    "ARLinkPredictor",
+    "InnerProductDecoder",
+    "GAE",
+    "VGAE",
+    "ARGA",
+    "ARGVA",
+    "DeepGraphInfomax",
+    "DeepGCNLayer",
+    "AttentiveFP",
+    "JumpingKnowledge",
+    "HeteroJumpingKnowledge",
+    "MaskLabel",
+    "MetaLayer",
+    "PMLP",
+    "Polynormer",
+    "BasicGNN",
+    "GCN",
+    "GraphSAGE",
+    "GIN",
+    "GAT",
+    "PNA",
+    "EdgeCNN",
+    "LabelPropagation",
+    "CorrectAndSmooth",
+    "LightGCN",
+    "BPRLoss",
+    "LINKX",
+    "SparseLinear",
+    "RECT_L",
+    "SignedGCN",
+    "NeuralFingerprint",
+    "GraphUNet",
+    "GroupAddRev",
+    "SGFormer",
+    "Node2Vec",
+    "MetaPath2Vec",
+    "RENet",
+    "TGNMemory",
+    "IdentityMessage",
+    "LastAggregator",
+    "MeanAggregator",
+    "TimeEncoder",
+    "LastNeighborLoader",
+    "SchNet",
+    "CFConv",
+    "SchNetInteractionBlock",
+    "GaussianSmearing",
+    "ShiftedSoftplus",
+    "RadiusInteractionGraph",
+    "DimeNet",
+    "DimeNetPlusPlus",
+    "BesselBasisLayer",
+    "SphericalBasisLayer",
+    "triplets",
+    "GNNFF",
+    "NodeBlock",
+    "EdgeBlock",
+    "GaussianFilter",
+    "GPSE",
+    "GPSENodeEncoder",
+    "GeneralLayer",
+    "GeneralMultiLayer",
+    "GNNStackStage",
+    "GNNInductiveHybridMultiHead",
+    "ViSNet",
+    "LPFormer",
+    "LPAttLayer",
+    "GraphMAE2",
+    "sce_loss",
+    "load_graphmae2_weights",
+    "download_graphmae2_checkpoint",
+    "Graphormer",
+    "GraphNodeFeature",
+    "GraphAttnBias",
+    "GraphormerMultiheadAttention",
+    "GraphormerGraphEncoderLayer",
+    "GraphormerGraphEncoder",
+    "load_graphormer_weights",
+    "download_graphormer_checkpoint",
+    "Graphormer3D",
+    "GaussianLayer",
+    "RBF",
+    "Graphormer3DEncoderLayer",
+    "NodeTaskHead",
+    "load_graphormer3d_weights",
+    "download_graphormer3d_checkpoint",
+    "to_captum_model",
+    "to_captum_input",
+    "captum_output_to_dicts",
+    "GPSModel",
+    "GPSLayer",
+    "CustomGatedGCN",
+    "AtomEncoder",
+    "BondEncoder",
+    "RWSEEncoder",
+    "SANGraphHead",
+    "load_gps_weights",
+    "download_gps_checkpoint",
+    "GROVER",
+    "GTransEncoder",
+    "Readout",
+    "load_grover_weights",
+    "download_grover_checkpoint",
+    "MoleBERT",
+    "MoleBERTGNN",
+    "MoleBERTGINConv",
+    "load_mole_bert_weights",
+    "download_mole_bert_checkpoint",
+    "UniMolModel",
+    "UniMolConfGenModel",
+    "UniMolDockingModel",
+    "UniMolGaussianLayer",
+    "UniMolNumericalEmbed",
+    "UniMolNonLinearHead",
+    "UniMolDistanceHead",
+    "UniMolClassificationHead",
+    "UniMolLinearHead",
+    "UniMolMaskLMHead",
+    "download_unimol_checkpoint",
+    "load_unimol_weights",
+    "UniMol2Model",
+    "UniMol2AtomFeature",
+    "UniMol2EdgeFeature",
+    "UniMol2SE3Kernel",
+    "UniMol2MovementHead",
+    "download_unimol2_checkpoint",
+    "load_unimol2_weights",
+    "UniMolPlusPCQModel",
+    "UniMolPlusOC20Model",
+    "UniMolPlusEnergyHead",
+    "download_unimol_plus_checkpoint",
+    "load_unimol_plus_weights",
+    "DockingPoseModelV2",
+    "download_unimol_docking_checkpoint",
+    "load_unimol_docking_weights",
+]
+
+# Inject Hugging Face Hub capabilities (from_pretrained, save_pretrained, push_to_hub, predict)
+# to all models in k3_node.models
+import keras
+from k3_node.hub.hub_mixin import K3NodeHubMixin
+
+def _defines_own(cls, attr):
+    """True if a k3_node class in ``cls``'s MRO defines ``attr`` itself (e.g. a model-specific
+    ``from_pretrained`` that loads original checkpoints), which must not be overwritten."""
+    return any(attr in vars(klass) for klass in cls.__mro__ if klass.__module__.startswith("k3_node"))
+
+
+def _is_graph_input(data):
+    if hasattr(data, "edge_index") or (hasattr(data, "z") and hasattr(data, "pos")):
+        return True
+    return isinstance(data, dict) and any(k in data for k in ("edge_index", "pos", "z"))
+
+
+def _graph_aware_predict(self, data=None, *args, **kwargs):
+    """Graph inputs (``Data``, ``Batch``, graph dicts) use the hub-style ``predict``; everything
+    else keeps Keras' batched ``Model.predict`` (arrays, ``tf.data``, ``PyDataset``, ...)."""
+    if _is_graph_input(data):
+        return K3NodeHubMixin.predict(self, data, *args, **kwargs)
+    return keras.Model.predict(self, data, *args, **kwargs)
+
+
+for _name in list(__all__):
+    _obj = globals().get(_name)
+    if isinstance(_obj, type) and issubclass(_obj, (keras.Model, keras.layers.Layer)):
+        if not issubclass(_obj, K3NodeHubMixin):
+            if not _defines_own(_obj, "from_pretrained"):
+                _obj.from_pretrained = classmethod(K3NodeHubMixin.from_pretrained.__func__)
+            if not _defines_own(_obj, "save_pretrained"):
+                _obj.save_pretrained = K3NodeHubMixin.save_pretrained
+            if not _defines_own(_obj, "push_to_hub"):
+                _obj.push_to_hub = K3NodeHubMixin.push_to_hub
+            if not _defines_own(_obj, "predict"):
+                if issubclass(_obj, keras.Model):
+                    _obj.predict = _graph_aware_predict
+                else:
+                    _obj.predict = K3NodeHubMixin.predict
+            if not hasattr(_obj, "_get_config"):
+                _obj._get_config = K3NodeHubMixin._get_config
+
+

k3_node/models/attentive_fp.py

@@ -0,0 +1,232 @@
+from typing import Optional
+import keras
+from keras import ops
+
+from k3_node.layers.conv import GATConv
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import softmax
+from k3_node.layers.pool import global_add_pool
+from k3_node.layers.pool.glob import _infer_size
+from k3_node.ops.segment import segment_sum
+
+try:
+    from keras.src.backend.common.symbolic_scope import in_symbolic_scope
+except ImportError:
+    def in_symbolic_scope():
+        return False
+
+
+def _gru_step(cell, x, h):
+    out, _ = cell(x, [h])
+    return out
+
+
+class GATEConv(MessagePassing):
+    r"""The edge-conditioned attention layer used as the first message
+    passing step of `AttentiveFP`."""
+    def __init__(self, in_channels: int, out_channels: int, edge_dim: int,
+                dropout: float = 0.0, **kwargs):
+        super().__init__(aggr="add", node_dim=0, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.edge_dim = edge_dim
+        self.dropout_rate = dropout
+
+        self.lin1 = keras.layers.Dense(out_channels, use_bias=False)
+        self.lin2 = keras.layers.Dense(out_channels, use_bias=False)
+        self.dropout = keras.layers.Dropout(dropout) if dropout > 0.0 else None
+
+        self.lin1.build((None, in_channels + edge_dim))
+        self.lin2.build((None, out_channels))
+
+        self.att_l = self.add_weight(shape=(1, out_channels), initializer="glorot_uniform", name="att_l")
+        self.att_r = self.add_weight(shape=(1, in_channels), initializer="glorot_uniform", name="att_r")
+        self.bias = self.add_weight(shape=(out_channels,), initializer="zeros", name="bias")
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def call(self, x, edge_index, edge_attr, training=None):
+        row, col = ops.cast(edge_index[0], "int32"), ops.cast(edge_index[1], "int32")
+
+        x_j = ops.take(x, row, axis=0)
+        x_i = ops.take(x, col, axis=0)
+
+        edge_attr = ops.cast(edge_attr, x.dtype)
+        h = ops.leaky_relu(self.lin1(ops.concatenate([x_j, edge_attr], axis=-1)), negative_slope=0.01)
+        alpha_j = ops.sum(h * self.att_l, axis=-1)
+        alpha_i = ops.sum(x_i * self.att_r, axis=-1)
+        alpha = ops.leaky_relu(alpha_j + alpha_i, negative_slope=0.01)
+
+        num_nodes = ops.shape(x)[0]
+        alpha = softmax(alpha, col, num_nodes=num_nodes, dim=0)
+        if self.dropout is not None:
+            alpha = self.dropout(alpha, training=training)
+
+        message = self.lin2(x_j) * ops.expand_dims(alpha, -1)
+        out = segment_sum(message, col, num_segments=num_nodes)
+        return out + self.bias
+
+
+class AttentiveFP(keras.Model):
+    r"""The Attentive FP model for molecular representation learning from the
+    `"Pushing the Boundaries of Molecular Representation for Drug Discovery
+    with the Graph Attention Mechanism"
+    <https://pubs.acs.org/doi/10.1021/acs.jmedchem.9b00959>`_ paper, based on
+    graph attention mechanisms.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Hidden node feature dimensionality.
+        out_channels (int): Size of each output sample.
+        edge_dim (int): Edge feature dimensionality.
+        num_layers (int): Number of GNN layers.
+        num_timesteps (int): Number of iterative refinement steps for global
+            readout.
+        dropout (float, optional): Dropout probability. (default: `0.0`)
+        batch_size (int, optional): Fixed batch size (number of graphs) for JAX/XLA
+            static shape compatibility. (default: `None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import AttentiveFP
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+        edge_attr = np.random.rand(30, 3).astype("float32")  # bond features
+
+        batch = np.repeat([0, 1], 5)  # two molecules with 5 atoms each
+        model = AttentiveFP(in_channels=8, hidden_channels=16, out_channels=1, edge_dim=3,
+                            num_layers=2, num_timesteps=2)
+        out = model(x, edge_index, edge_attr, batch)  # one prediction per molecule
+        print(tuple(out.shape))  # (2, 1)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        edge_dim: int,
+        num_layers: int,
+        num_timesteps: int,
+        dropout: float = 0.0,
+        batch_size: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.edge_dim = edge_dim
+        self.num_layers = num_layers
+        self.num_timesteps = num_timesteps
+        self.dropout_rate = dropout
+        self.batch_size = batch_size
+
+        self.lin1 = keras.layers.Dense(hidden_channels)
+        self.lin1.build((None, in_channels))
+
+        self.gate_conv = GATEConv(hidden_channels, hidden_channels, edge_dim, dropout)
+        self.gru = keras.layers.GRUCell(hidden_channels)
+        self.gru.build((None, hidden_channels))
+
+        self.atom_convs = []
+        self.atom_grus = []
+        for _ in range(num_layers - 1):
+            conv = GATConv(hidden_channels, hidden_channels, dropout=dropout,
+                           add_self_loops=False, negative_slope=0.01)
+            conv.build((None, hidden_channels))
+            self.atom_convs.append(conv)
+            gru = keras.layers.GRUCell(hidden_channels)
+            gru.build((None, hidden_channels))
+            self.atom_grus.append(gru)
+
+        self.mol_conv = GATConv(hidden_channels, hidden_channels, dropout=dropout,
+                                add_self_loops=False, negative_slope=0.01)
+        self.mol_conv.build([(None, hidden_channels), (None, hidden_channels)])
+        self.mol_gru = keras.layers.GRUCell(hidden_channels)
+        self.mol_gru.build((None, hidden_channels))
+
+        self.lin2 = keras.layers.Dense(out_channels)
+        self.lin2.build((None, hidden_channels))
+
+        self.dropout = keras.layers.Dropout(dropout) if dropout > 0.0 else None
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_attr=None, batch=None, batch_size=None, training=None):
+        if isinstance(x, dict):
+            edge_index = x.get("edge_index")
+            edge_attr = x.get("edge_attr")
+            batch = x.get("batch")
+            batch_size = x.get("batch_size", batch_size)
+            x = x.get("x")
+        elif isinstance(x, (tuple, list)) and edge_index is None:
+            if len(x) >= 4:
+                x, edge_index, edge_attr, batch = x[0], x[1], x[2], x[3]
+            elif len(x) == 3:
+                x, edge_index, edge_attr = x[0], x[1], x[2]
+
+        bs = batch_size if batch_size is not None else self.batch_size
+        x = ops.cast(x, "float32")
+        if edge_attr is not None:
+            edge_attr = ops.cast(edge_attr, "float32")
+        # Atom Embedding:
+        x = ops.leaky_relu(self.lin1(x), negative_slope=0.01)
+
+        h = ops.elu(self.gate_conv(x, edge_index, edge_attr, training=training))
+        if self.dropout is not None:
+            h = self.dropout(h, training=training)
+        x = ops.relu(_gru_step(self.gru, h, x))
+
+        for conv, gru in zip(self.atom_convs, self.atom_grus):
+            h = conv(x, edge_index, training=training)
+            h = ops.elu(h)
+            if self.dropout is not None:
+                h = self.dropout(h, training=training)
+            x = ops.relu(_gru_step(gru, h, x))
+
+        # Molecule Embedding:
+        if batch is None:
+            batch = ops.zeros((ops.shape(x)[0],), dtype="int32")
+        else:
+            batch = ops.cast(batch, "int32")
+        num_nodes = ops.shape(batch)[0]
+        row = ops.arange(num_nodes, dtype="int32")
+        mol_edge_index = ops.stack([row, batch], axis=0)
+
+        if keras.config.backend() == "jax":
+            size = bs
+        elif in_symbolic_scope():
+            size = bs
+        else:
+            size = _infer_size(batch)
+            if size is None:
+                size = bs
+
+        out = ops.relu(global_add_pool(x, batch, size=size))
+        for _ in range(self.num_timesteps):
+            h = ops.elu(self.mol_conv((x, out), mol_edge_index, training=training))
+            if self.dropout is not None:
+                h = self.dropout(h, training=training)
+            out = ops.relu(_gru_step(self.mol_gru, h, out))
+
+        # Predictor:
+        if self.dropout is not None:
+            out = self.dropout(out, training=training)
+        return self.lin2(out)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"in_channels={self.in_channels}, "
+            f"hidden_channels={self.hidden_channels}, "
+            f"out_channels={self.out_channels}, "
+            f"edge_dim={self.edge_dim}, "
+            f"num_layers={self.num_layers}, "
+            f"num_timesteps={self.num_timesteps}"
+            f")"
+        )