k3-node 1.0.0__py3-none-any.whl

k3_node/__init__.py

@@ -0,0 +1,122 @@
+"""
+`k3_node` is a library for building multibackend graph neural networks. 
+Built upon Keras 3.0 the models can be trained using TensorFlow, PyTorch, 
+or JAX.
+
+To install the package, run:
+    
+```bash
+git clone https://github.com/anas-rz/k3-node.git # bash
+```
+        
+```python
+# in your code
+import sys
+sys.path.append('k3-node')
+
+import os
+os.environ['KERAS_BACKEND'] = 'tensorflow' # or 'torch' or 'jax'
+
+from k3_node import ...
+```
+"""
+__version__ = "1.0.0"
+
+import warnings
+
+# Index tensors are requested as int64; without `jax_enable_x64`, JAX stores them as int32
+# (which is all graphs of this size need) and would warn on every conversion.
+warnings.filterwarnings("ignore", message="Explicitly requested dtype int64", category=UserWarning)
+
+from k3_node import data
+from k3_node.data import Data, Batch
+from k3_node import datasets
+from k3_node import io
+from k3_node import layers
+from k3_node import loader
+from k3_node import transforms
+from k3_node import models
+from k3_node import applications
+
+from k3_node import metrics
+from k3_node import tasks
+from k3_node.tasks import (
+    NodeClassifier,
+    NodeRegressor,
+    GraphClassifier,
+    GraphRegressor,
+    LinkPredictor,
+)
+
+from k3_node import etl
+from k3_node.etl import (
+    TableToGraph,
+    TabularToGraph,
+    table_to_graph,
+    RelationalToGraph,
+    relational_to_graph,
+)
+
+from k3_node import hub
+from k3_node.hub import (
+    from_pretrained,
+    push_to_hub,
+    save_pretrained,
+    load_dataset_from_hub,
+    push_dataset_to_hub,
+)
+
+from k3_node import export
+from k3_node.export import (
+    export_onnx,
+    export_tflite,
+    export_tensorrt,
+    generate_triton_config,
+    ONNXModel,
+    TFLiteModel,
+)
+
+from k3_node import rag
+from k3_node.layers import kge
+
+__all__ = [
+    "data",
+    "Data",
+    "Batch",
+    "datasets",
+    "io",
+    "layers",
+    "loader",
+    "transforms",
+    "models",
+    "applications",
+    "metrics",
+    "tasks",
+    "NodeClassifier",
+    "NodeRegressor",
+    "GraphClassifier",
+    "GraphRegressor",
+    "LinkPredictor",
+    "etl",
+    "TableToGraph",
+    "TabularToGraph",
+    "table_to_graph",
+    "RelationalToGraph",
+    "relational_to_graph",
+    "hub",
+    "from_pretrained",
+    "push_to_hub",
+    "save_pretrained",
+    "load_dataset_from_hub",
+    "push_dataset_to_hub",
+    "export",
+    "export_onnx",
+    "export_tflite",
+    "export_tensorrt",
+    "generate_triton_config",
+    "ONNXModel",
+    "TFLiteModel",
+    "rag",
+    "kge",
+]
+

k3_node/applications/__init__.py

@@ -0,0 +1,17 @@
+"""Domain-specific applications built on K3-Node.
+
+Submodules:
+- `chemistry`: Molecular graphs, quantum property prediction, SMILES processing, Uni-Mol, AttentiveFP, SchNet, DimeNet.
+- `bio`: Macromolecular structures, protein-protein interactions, docking, UniMolDocking.
+- `materials`: Crystal graph neural networks, periodic boundaries, CHGNet, M3GNet, MEGNet, TensorNet, SO3Net, QET.
+"""
+
+from k3_node.applications import bio
+from k3_node.applications import chemistry
+from k3_node.applications import materials
+
+__all__ = [
+    "bio",
+    "chemistry",
+    "materials",
+]

k3_node/applications/bio/__init__.py

@@ -0,0 +1,21 @@
+"""Biological and macromolecular graph neural network models and applications."""
+
+from k3_node.models.unimol import (
+    UniMolDockingModel,
+    download_unimol_checkpoint,
+    load_unimol_weights,
+)
+from k3_node.models.unimol_docking_v2 import (
+    DockingPoseModelV2,
+    download_unimol_docking_checkpoint,
+    load_unimol_docking_weights,
+)
+
+__all__ = [
+    "UniMolDockingModel",
+    "DockingPoseModelV2",
+    "download_unimol_checkpoint",
+    "load_unimol_weights",
+    "download_unimol_docking_checkpoint",
+    "load_unimol_docking_weights",
+]

k3_node/applications/chemistry/__init__.py

@@ -0,0 +1,155 @@
+"""Chemistry and molecular graph neural network models and applications."""
+
+from k3_node.models.attentive_fp import AttentiveFP
+from k3_node.models.dimenet import (
+    DimeNet,
+    DimeNetPlusPlus,
+    BesselBasisLayer,
+    SphericalBasisLayer,
+    triplets,
+)
+from k3_node.models.grover import (
+    GROVER,
+    GTransEncoder,
+    Readout as GROVERReadout,
+    load_grover_weights,
+    download_grover_checkpoint,
+)
+from k3_node.models.mole_bert import (
+    MoleBERT,
+    load_mole_bert_weights,
+    download_mole_bert_checkpoint,
+)
+from k3_node.models.neural_fingerprint import NeuralFingerprint
+from k3_node.models.schnet import (
+    SchNet,
+    CFConv,
+    GaussianSmearing,
+    ShiftedSoftplus,
+    RadiusInteractionGraph,
+)
+from k3_node.models.visnet import ViSNet
+from k3_node.models.gnnff import (
+    GNNFF,
+    NodeBlock,
+    EdgeBlock,
+    GaussianFilter,
+)
+from k3_node.models.graphormer import (
+    Graphormer,
+    GraphNodeFeature,
+    GraphAttnBias,
+    GraphormerMultiheadAttention,
+    GraphormerGraphEncoderLayer,
+    GraphormerGraphEncoder,
+    load_graphormer_weights,
+    download_graphormer_checkpoint,
+)
+from k3_node.models.graphormer_3d import (
+    Graphormer3D,
+    GaussianLayer as Graphormer3DGaussianLayer,
+    RBF,
+    Graphormer3DEncoderLayer,
+    NodeTaskHead,
+    load_graphormer3d_weights,
+    download_graphormer3d_checkpoint,
+)
+from k3_node.models.unimol import (
+    UniMolModel,
+    UniMolConfGenModel,
+    GaussianLayer,
+    NumericalEmbed,
+    NonLinearHead,
+    DistanceHead,
+    ClassificationHead,
+    LinearHead,
+    MaskLMHead,
+    download_unimol_checkpoint,
+    load_unimol_weights,
+)
+from k3_node.models.unimol2 import (
+    UniMol2Model,
+    AtomFeature,
+    EdgeFeature,
+    SE3InvariantKernel,
+    MovementPredictionHead,
+    UniMol2TransformerLayer,
+    download_unimol2_checkpoint,
+    load_unimol2_weights,
+)
+from k3_node.models.unimol_plus import (
+    UniMolPlusPCQModel,
+    UniMolPlusOC20Model,
+    UnimolPlusEncoderLayer,
+    EnergyHead,
+    download_unimol_plus_checkpoint,
+    load_unimol_plus_weights,
+)
+
+__all__ = [
+    "AttentiveFP",
+    "DimeNet",
+    "DimeNetPlusPlus",
+    "BesselBasisLayer",
+    "SphericalBasisLayer",
+    "triplets",
+    "GROVER",
+    "GTransEncoder",
+    "GROVERReadout",
+    "load_grover_weights",
+    "download_grover_checkpoint",
+    "MoleBERT",
+    "load_mole_bert_weights",
+    "download_mole_bert_checkpoint",
+    "NeuralFingerprint",
+    "SchNet",
+    "CFConv",
+    "GaussianSmearing",
+    "ShiftedSoftplus",
+    "RadiusInteractionGraph",
+    "ViSNet",
+    "GNNFF",
+    "NodeBlock",
+    "EdgeBlock",
+    "GaussianFilter",
+    "Graphormer",
+    "GraphNodeFeature",
+    "GraphAttnBias",
+    "GraphormerMultiheadAttention",
+    "GraphormerGraphEncoderLayer",
+    "GraphormerGraphEncoder",
+    "load_graphormer_weights",
+    "download_graphormer_checkpoint",
+    "Graphormer3D",
+    "Graphormer3DGaussianLayer",
+    "RBF",
+    "Graphormer3DEncoderLayer",
+    "NodeTaskHead",
+    "load_graphormer3d_weights",
+    "download_graphormer3d_checkpoint",
+    "UniMolModel",
+    "UniMolConfGenModel",
+    "GaussianLayer",
+    "NumericalEmbed",
+    "NonLinearHead",
+    "DistanceHead",
+    "ClassificationHead",
+    "LinearHead",
+    "MaskLMHead",
+    "download_unimol_checkpoint",
+    "load_unimol_weights",
+    "UniMol2Model",
+    "AtomFeature",
+    "EdgeFeature",
+    "SE3InvariantKernel",
+    "MovementPredictionHead",
+    "UniMol2TransformerLayer",
+    "download_unimol2_checkpoint",
+    "load_unimol2_weights",
+    "UniMolPlusPCQModel",
+    "UniMolPlusOC20Model",
+    "UnimolPlusEncoderLayer",
+    "EnergyHead",
+    "download_unimol_plus_checkpoint",
+    "load_unimol_plus_weights",
+]

k3_node/applications/materials/__init__.py

@@ -0,0 +1,127 @@
+"""Materials and crystal graph neural network models (multi-backend Keras 3 ports of MatGL)."""
+
+from .core import (
+    SoftPlus2,
+    SoftExponential,
+    MLP,
+    GatedMLP,
+    EmbeddingBlock,
+    vector_to_skewtensor,
+    vector_to_symtensor,
+    decompose_tensor,
+    new_radial_tensor,
+    tensor_norm,
+)
+from .basis import (
+    GaussianExpansion,
+    RadialBesselFunction,
+    FourierExpansion,
+    ChebyshevRadialBasis,
+    SphericalBesselFunction,
+    SphericalBesselWithHarmonics,
+    BondExpansion,
+    polynomial_cutoff,
+    cosine_cutoff,
+    compute_pair_vector_and_distance,
+    compute_theta,
+    compute_theta_and_phi,
+)
+from .readout import (
+    ReduceReadOut,
+    WeightedReadOut,
+    WeightedAtomReadOut,
+    Set2SetReadOut,
+    EdgeSet2Set,
+)
+from .megnet import MEGNet, MEGNetBlock, MEGNetGraphConv
+from .m3gnet import M3GNet, M3GNetBlock, M3GNetGraphConv, ThreeBodyInteractions
+from .tensornet import TensorNet, TensorEmbedding, TensorNetInteraction
+from .chgnet import CHGNet, CHGNetAtomGraphBlock, CHGNetBondGraphBlock
+from .so3net import SO3Net, SO3Convolution, RealSphericalHarmonics
+from .grace import GRACE, GraceSPBasis, GraceACEStack
+from .qet import QET, LinearQeq, ElectrostaticPotential
+from .wrappers import TransformedTargetModel, Potential
+from .io import (
+    download_matgl_checkpoint,
+    load_matgl_weights,
+    load_model,
+    get_available_pretrained_models,
+)
+
+__all__ = [
+    # Models
+    "MEGNet",
+    "MEGNetBlock",
+    "MEGNetGraphConv",
+    "M3GNet",
+    "M3GNetBlock",
+    "M3GNetGraphConv",
+    "ThreeBodyInteractions",
+    "TensorNet",
+    "TensorEmbedding",
+    "TensorNetInteraction",
+    "CHGNet",
+    "CHGNetAtomGraphBlock",
+    "CHGNetBondGraphBlock",
+    "SO3Net",
+    "SO3Convolution",
+    "RealSphericalHarmonics",
+    "GRACE",
+    "GraceSPBasis",
+    "GraceACEStack",
+    "QET",
+    "LinearQeq",
+    "ElectrostaticPotential",
+    "TransformedTargetModel",
+    "Potential",
+    # Basis & Geometry
+    "GaussianExpansion",
+    "RadialBesselFunction",
+    "FourierExpansion",
+    "ChebyshevRadialBasis",
+    "SphericalBesselFunction",
+    "SphericalBesselWithHarmonics",
+    "BondExpansion",
+    "polynomial_cutoff",
+    "cosine_cutoff",
+    "compute_pair_vector_and_distance",
+    "compute_theta",
+    "compute_theta_and_phi",
+    # Core & Readout
+    "SoftPlus2",
+    "SoftExponential",
+    "MLP",
+    "GatedMLP",
+    "EmbeddingBlock",
+    "vector_to_skewtensor",
+    "vector_to_symtensor",
+    "decompose_tensor",
+    "new_radial_tensor",
+    "tensor_norm",
+    "ReduceReadOut",
+    "WeightedReadOut",
+    "WeightedAtomReadOut",
+    "Set2SetReadOut",
+    "EdgeSet2Set",
+    # Checkpoint I/O
+    "download_matgl_checkpoint",
+    "load_matgl_weights",
+    "load_model",
+    "get_available_pretrained_models",
+]
+
+# Inject Hugging Face Hub capabilities
+import keras
+from k3_node.hub.hub_mixin import K3NodeHubMixin
+
+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):
+            _obj.from_pretrained = classmethod(K3NodeHubMixin.from_pretrained.__func__)
+            _obj.save_pretrained = K3NodeHubMixin.save_pretrained
+            _obj.push_to_hub = K3NodeHubMixin.push_to_hub
+            _obj.predict = K3NodeHubMixin.predict
+            if not hasattr(_obj, "_get_config"):
+                _obj._get_config = K3NodeHubMixin._get_config
+