k3-node 1.0.0__py3-none-any.whl

k3_node/models/test_correct_and_smooth.py

@@ -0,0 +1,40 @@
+from keras import ops
+
+from k3_node.models import CorrectAndSmooth
+
+
+def test_correct_and_smooth():
+    y_soft = ops.repeat(ops.convert_to_tensor([[0.1, 0.5, 0.4]]), 6, axis=0)
+    y_true = ops.convert_to_tensor([1, 0, 0, 2, 1, 1])
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2, 4, 5], [1, 0, 2, 1, 5, 4]])
+    mask = ops.convert_to_tensor([True, False, True, False, True, False])
+
+    model = CorrectAndSmooth(
+        num_correction_layers=2,
+        correction_alpha=0.5,
+        num_smoothing_layers=2,
+        smoothing_alpha=0.5,
+    )
+    assert str(model) == ('CorrectAndSmooth(\n'
+                          '  correct: num_layers=2, alpha=0.5\n'
+                          '  smooth:  num_layers=2, alpha=0.5\n'
+                          '  autoscale=True, scale=1.0\n'
+                          ')')
+
+    out = model.correct(y_soft, y_true[mask], mask, edge_index)
+    assert ops.shape(out) == (6, 3)
+
+    out = model.smooth(y_soft, y_true[mask], mask, edge_index)
+    assert ops.shape(out) == (6, 3)
+
+    # Without autoscale:
+    model_no_auto = CorrectAndSmooth(
+        num_correction_layers=2,
+        correction_alpha=0.5,
+        num_smoothing_layers=2,
+        smoothing_alpha=0.5,
+        autoscale=False,
+    )
+    out = model_no_auto.correct(y_soft, y_true[mask], mask, edge_index)
+    assert ops.shape(out) == (6, 3)
+

k3_node/models/test_deep_graph_infomax.py

@@ -0,0 +1,68 @@
+import numpy as np
+from keras import ops
+
+from k3_node.models import DeepGraphInfomax
+
+
+def test_infomax():
+    model = DeepGraphInfomax(
+        hidden_channels=16,
+        encoder=lambda x: x,
+        summary=lambda z, *args: ops.mean(z, axis=0),
+        corruption=lambda x: x + 1,
+    )
+    assert str(model) == "DeepGraphInfomax(16)"
+
+    x = ops.ones((20, 16))
+
+    pos_z, neg_z, summary = model(x)
+    assert ops.shape(pos_z) == (20, 16)
+    assert ops.shape(neg_z) == (20, 16)
+    assert ops.shape(summary) == (16,)
+
+    loss = model.loss(pos_z, neg_z, summary)
+    assert float(ops.convert_to_numpy(loss)) >= 0
+
+    acc = model.test(
+        train_z=ops.ones((20, 16)),
+        train_y=ops.convert_to_tensor(np.random.randint(0, 10, (20,))),
+        test_z=ops.ones((20, 16)),
+        test_y=ops.convert_to_tensor(np.random.randint(0, 10, (20,))),
+    )
+    assert 0 <= acc <= 1
+
+
+def test_infomax_predefined_model():
+    from k3_node.layers.conv import GCNConv
+
+    class Encoder:
+        def __init__(self):
+            self.conv1 = GCNConv(16, 16)
+            self.conv2 = GCNConv(16, 16)
+
+        def __call__(self, x, edge_index, edge_weight=None):
+            x = ops.relu(self.conv1(x, edge_index, edge_weight=edge_weight))
+            return self.conv2(x, edge_index, edge_weight=edge_weight)
+
+    def corruption(x, edge_index, edge_weight):
+        perm = np.random.permutation(ops.shape(x)[0])
+        return ops.take(x, ops.convert_to_tensor(perm), axis=0), edge_index, edge_weight
+
+    model = DeepGraphInfomax(
+        hidden_channels=16,
+        encoder=Encoder(),
+        summary=lambda z, *args, **kwargs: ops.sigmoid(ops.mean(z, axis=0)),
+        corruption=corruption,
+    )
+
+    x = ops.convert_to_tensor(np.random.randn(4, 16).astype("float32"))
+    edge_index = ops.convert_to_tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]], dtype="int64")
+    edge_weight = ops.convert_to_tensor(np.random.rand(edge_index.shape[1]).astype("float32"))
+
+    pos_z, neg_z, summary = model(x, edge_index, edge_weight=edge_weight)
+    assert ops.shape(pos_z) == (4, 16)
+    assert ops.shape(neg_z) == (4, 16)
+    assert ops.shape(summary) == (16,)
+
+    loss = model.loss(pos_z, neg_z, summary)
+    assert float(ops.convert_to_numpy(loss)) >= 0

k3_node/models/test_deepgcn.py

@@ -0,0 +1,21 @@
+import pytest
+from keras import ops
+
+from k3_node.layers.conv import GENConv
+from k3_node.layers.norm import LayerNorm
+from k3_node.models import DeepGCNLayer
+
+
+@pytest.mark.parametrize("block_tuple", [("res+", 1), ("res", 1), ("dense", 2), ("plain", 1)])
+def test_deepgcn(block_tuple):
+    block, expansion = block_tuple
+    x = ops.convert_to_tensor([[1.0] * 8] * 3, dtype="float32")
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype="int64")
+    conv = GENConv(8, 8)
+    norm = LayerNorm(8)
+    act = ops.relu
+    layer = DeepGCNLayer(conv, norm, act, block=block)
+    assert str(layer) == f"DeepGCNLayer(block={block})"
+
+    out = layer(x, edge_index)
+    assert ops.shape(out) == (3, 8 * expansion)

k3_node/models/test_dimenet.py

@@ -0,0 +1,86 @@
+import keras.ops as ops
+from k3_node.models.dimenet import (
+    DimeNet,
+    DimeNetPlusPlus,
+    BesselBasisLayer,
+    SphericalBasisLayer,
+    triplets,
+)
+
+
+def test_triplets():
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 1, 2, 0, 2],
+        [1, 0, 2, 1, 2, 0],
+    ])
+    col, row, idx_i, idx_j, idx_k, idx_kj, idx_ji = triplets(edge_index, num_nodes=3)
+    assert len(idx_i) > 0
+    assert len(idx_kj) == len(idx_ji)
+
+
+def test_bessel_basis_layer():
+    bessel = BesselBasisLayer(num_radial=6, cutoff=5.0)
+    dist = ops.convert_to_tensor([1.0, 2.0, 3.0])
+    out = bessel(dist)
+    assert out.shape == (3, 6)
+
+
+def test_spherical_basis_layer():
+    sbf = SphericalBasisLayer(num_spherical=3, num_radial=6, cutoff=5.0)
+    dist = ops.convert_to_tensor([1.0, 2.0])
+    angle = ops.convert_to_tensor([0.5, 1.2])
+    idx_kj = ops.convert_to_tensor([0, 1])
+    out = sbf(dist, angle, idx_kj)
+    assert out.shape == (2, 18)
+
+
+def test_dimenet():
+    z = ops.convert_to_tensor([1, 6, 8, 1])
+    pos = ops.convert_to_tensor([
+        [0.0, 0.0, 0.0],
+        [1.0, 0.0, 0.0],
+        [0.0, 1.0, 0.0],
+        [1.0, 1.0, 0.0],
+    ], dtype="float32")
+
+    model = DimeNet(
+        hidden_channels=16,
+        out_channels=1,
+        num_blocks=2,
+        num_bilinear=8,
+        num_spherical=3,
+        num_radial=6,
+        cutoff=5.0,
+    )
+    out = model(z, pos)
+    assert out.shape == (1,)
+
+    # With batch
+    batch = ops.convert_to_tensor([0, 0, 1, 1])
+    out_b = model(z, pos, batch=batch)
+    assert out_b.shape == (2, 1)
+
+
+def test_dimenet_plus_plus():
+    z = ops.convert_to_tensor([1, 6, 8, 1])
+    pos = ops.convert_to_tensor([
+        [0.0, 0.0, 0.0],
+        [1.0, 0.0, 0.0],
+        [0.0, 1.0, 0.0],
+        [1.0, 1.0, 0.0],
+    ], dtype="float32")
+
+    model = DimeNetPlusPlus(
+        hidden_channels=16,
+        out_channels=1,
+        num_blocks=2,
+        int_emb_size=8,
+        basis_emb_size=8,
+        out_emb_channels=16,
+        num_spherical=3,
+        num_radial=6,
+        cutoff=5.0,
+    )
+    out = model(z, pos)
+    assert out.shape == (1,)
+

k3_node/models/test_domain_apis.py

@@ -0,0 +1,138 @@
+"""Test domain-specific API organization (applications: materials, bio, chemistry)."""
+
+import pytest
+
+
+def test_clean_main_api():
+    """Ensure main k3_node API is clean and domain packages are under applications."""
+    import k3_node
+
+    assert not hasattr(k3_node, "materials"), "k3_node should not expose materials directly"
+    assert not hasattr(k3_node, "bio"), "k3_node should not expose bio directly"
+    assert not hasattr(k3_node, "chemistry"), "k3_node should not expose chemistry directly"
+
+    assert hasattr(k3_node, "applications"), "k3_node must expose applications"
+    assert hasattr(k3_node.applications, "materials"), "applications must expose materials"
+    assert hasattr(k3_node.applications, "bio"), "applications must expose bio"
+    assert hasattr(k3_node.applications, "chemistry"), "applications must expose chemistry"
+
+
+def test_materials_application_api():
+    # Via k3_node.applications.materials
+    from k3_node.applications.materials import (
+        MEGNet,
+        M3GNet,
+        TensorNet,
+        CHGNet,
+        SO3Net,
+        GRACE,
+        QET,
+        TransformedTargetModel,
+        Potential,
+        download_matgl_checkpoint,
+        load_matgl_weights,
+        load_model,
+        get_available_pretrained_models,
+        basis,
+        core,
+        readout,
+    )
+    assert MEGNet is not None
+    assert M3GNet is not None
+    assert TensorNet is not None
+    assert CHGNet is not None
+    assert SO3Net is not None
+    assert GRACE is not None
+    assert QET is not None
+    assert basis is not None
+    assert core is not None
+    assert readout is not None
+
+    # Via k3_node.models.materials (backward compatibility)
+    from k3_node.models.materials import (
+        MEGNet as MatMEGNet,
+        M3GNet as MatM3GNet,
+        TensorNet as MatTensorNet,
+    )
+    assert MatMEGNet is MEGNet
+    assert MatM3GNet is M3GNet
+    assert MatTensorNet is TensorNet
+
+
+def test_bio_application_api():
+    # Via k3_node.applications.bio
+    from k3_node.applications.bio import (
+        UniMolDockingModel,
+        DockingPoseModelV2,
+        download_unimol_checkpoint,
+        load_unimol_weights,
+        download_unimol_docking_checkpoint,
+        load_unimol_docking_weights,
+    )
+    assert UniMolDockingModel is not None
+    assert DockingPoseModelV2 is not None
+
+    # Via k3_node.models.bio (backward compatibility)
+    from k3_node.models.bio import (
+        UniMolDockingModel as BioDocking1,
+        DockingPoseModelV2 as BioDocking2,
+    )
+    assert BioDocking1 is UniMolDockingModel
+    assert BioDocking2 is DockingPoseModelV2
+
+
+def test_chemistry_application_api():
+    # Via k3_node.applications.chemistry
+    from k3_node.applications.chemistry import (
+        AttentiveFP,
+        DimeNet,
+        DimeNetPlusPlus,
+        GROVER,
+        MoleBERT,
+        NeuralFingerprint,
+        SchNet,
+        ViSNet,
+        GNNFF,
+        Graphormer,
+        Graphormer3D,
+        UniMolModel,
+        UniMolConfGenModel,
+        UniMol2Model,
+        UniMolPlusPCQModel,
+        UniMolPlusOC20Model,
+    )
+    assert AttentiveFP is not None
+    assert DimeNet is not None
+    assert SchNet is not None
+    assert UniMolModel is not None
+
+    # Via k3_node.models.chemistry (backward compatibility)
+    from k3_node.models.chemistry import (
+        SchNet as ChemSchNet,
+        DimeNet as ChemDimeNet,
+        UniMolModel as ChemUniMol,
+    )
+    assert ChemSchNet is SchNet
+    assert ChemDimeNet is DimeNet
+    assert ChemUniMol is UniMolModel
+
+
+def test_backward_compatibility():
+    # All models still exportable from top-level k3_node.models
+    from k3_node.models import (
+        MEGNet,
+        M3GNet,
+        TensorNet,
+        CHGNet,
+        SO3Net,
+        GRACE,
+        QET,
+        SchNet,
+        DimeNet,
+        UniMolModel,
+        UniMolDockingModel,
+        DockingPoseModelV2,
+    )
+    assert MEGNet is not None
+    assert SchNet is not None
+    assert UniMolModel is not None

k3_node/models/test_gnnff.py

@@ -0,0 +1,24 @@
+import keras.ops as ops
+from k3_node.models.gnnff import GNNFF, GaussianFilter
+
+
+def test_gaussian_filter():
+    gf = GaussianFilter(start=0.0, stop=5.0, num_gaussians=10)
+    dist = ops.convert_to_tensor([0.5, 1.5, 3.0])
+    out = gf(dist)
+    assert out.shape == (3, 10)
+
+
+def test_gnnff():
+    z = ops.convert_to_tensor([1, 6, 8, 1])
+    pos = ops.convert_to_tensor([
+        [0.0, 0.0, 0.0],
+        [1.0, 0.0, 0.0],
+        [0.0, 1.0, 0.0],
+        [1.0, 1.0, 0.0],
+    ], dtype="float32")
+
+    model = GNNFF(hidden_node_channels=16, hidden_edge_channels=16, num_layers=2)
+    force = model(z, pos)
+    assert force.shape == (4, 3)
+

k3_node/models/test_gps_model.py

@@ -0,0 +1,271 @@
+import os
+import tempfile
+import numpy as np
+import pytest
+import keras
+from keras import ops
+
+from k3_node.models.gps_model import (
+    AtomEncoder,
+    BondEncoder,
+    RWSEEncoder,
+    CustomGatedGCN,
+    GPSLayer,
+    SANGraphHead,
+    GPSModel,
+    load_gps_weights,
+    download_gps_checkpoint,
+)
+
+
+def test_atom_and_bond_encoders():
+    # 9 features for atoms, 3 for bonds
+    x = ops.convert_to_tensor(np.array([[6, 0, 4, 4, 3, 2, 2, 0, 0],
+                                        [8, 0, 2, 2, 2, 1, 1, 0, 0],
+                                        [1, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=np.int32))
+    atom_enc = AtomEncoder(emb_dim=64)
+    x_emb = atom_enc(x)
+    assert ops.shape(x_emb) == (3, 64)
+
+    edge_attr = ops.convert_to_tensor(np.array([[0, 0, 0],
+                                                [1, 0, 0]], dtype=np.int32))
+    bond_enc = BondEncoder(emb_dim=64)
+    e_emb = bond_enc(edge_attr)
+    assert ops.shape(e_emb) == (2, 64)
+
+
+def test_rwse_encoder():
+    pestat = ops.convert_to_tensor(np.ones((4, 16), dtype=np.float32))
+    rwse = RWSEEncoder(num_rw_steps=16, pe_dim=20)
+    pe_out = rwse(pestat, training=False)
+    assert ops.shape(pe_out) == (4, 20)
+
+
+def test_custom_gated_gcn():
+    x = ops.convert_to_tensor(np.random.randn(4, 32).astype(np.float32))
+    e = ops.convert_to_tensor(np.random.randn(6, 32).astype(np.float32))
+    edge_index = ops.convert_to_tensor(np.array([[0, 1, 1, 2, 2, 3],
+                                                 [1, 0, 2, 1, 3, 2]], dtype=np.int32))
+
+    layer = CustomGatedGCN(in_dim=32, out_dim=32, dropout=0.0, residual=True, act="gelu")
+    x_out, e_out = layer(x, edge_index, e, training=False)
+
+    assert ops.shape(x_out) == (4, 32)
+    assert ops.shape(e_out) == (6, 32)
+
+
+def test_gps_layer():
+    dim_h = 32
+    num_heads = 4
+    x = ops.convert_to_tensor(np.random.randn(5, dim_h).astype(np.float32))
+    e = ops.convert_to_tensor(np.random.randn(6, dim_h).astype(np.float32))
+    edge_index = ops.convert_to_tensor(np.array([[0, 1, 1, 2, 3, 4],
+                                                 [1, 0, 2, 1, 4, 3]], dtype=np.int32))
+    # Two graphs: graph 0 has 3 nodes, graph 1 has 2 nodes
+    batch = ops.convert_to_tensor(np.array([0, 0, 0, 1, 1], dtype=np.int32))
+
+    layer = GPSLayer(
+        dim_h=dim_h,
+        local_gnn_type="CustomGatedGCN",
+        global_model_type="Transformer",
+        num_heads=num_heads,
+        act="gelu",
+        dropout=0.0,
+        attn_dropout=0.0,
+        batch_norm=True,
+    )
+
+    x_out, e_out = layer(x, edge_index, e, batch=batch, training=False)
+    assert ops.shape(x_out) == (5, dim_h)
+    assert ops.shape(e_out) == (6, dim_h)
+
+
+def test_san_graph_head():
+    x = ops.convert_to_tensor(np.random.randn(6, 64).astype(np.float32))
+    batch = ops.convert_to_tensor(np.array([0, 0, 0, 1, 1, 1], dtype=np.int32))
+
+    head = SANGraphHead(dim_in=64, dim_out=1, L=2, act="gelu", pooling="mean")
+    pred = head(x, batch=batch, training=False)
+    assert ops.shape(pred) == (2, 1)
+
+    head_sum = SANGraphHead(dim_in=64, dim_out=2, L=1, act="relu", pooling="sum")
+    pred_sum = head_sum(x, batch=batch, training=False)
+    assert ops.shape(pred_sum) == (2, 2)
+
+
+def test_gps_model_forward():
+    model = GPSModel(
+        dim_in=32,
+        dim_out=1,
+        num_layers=2,
+        dim_hidden=32,
+        num_heads=4,
+        local_gnn_type="CustomGatedGCN",
+        act="gelu",
+        dropout=0.1,
+        attn_dropout=0.1,
+        batch_norm=True,
+        node_encoder_type="Atom+RWSE",
+        edge_encoder_type="Bond",
+        rwse_num_steps=16,
+        rwse_dim_pe=8,
+        graph_pooling="mean",
+        head_layers=2,
+    )
+
+    x = ops.convert_to_tensor(np.array([[6, 0, 4, 4, 3, 2, 2, 0, 0],
+                                        [8, 0, 2, 2, 2, 1, 1, 0, 0],
+                                        [6, 0, 3, 3, 3, 2, 2, 0, 0]], dtype=np.int32))
+    edge_index = ops.convert_to_tensor(np.array([[0, 1, 1, 2],
+                                                 [1, 0, 2, 1]], dtype=np.int32))
+    edge_attr = ops.convert_to_tensor(np.array([[0, 0, 0],
+                                                [0, 0, 0],
+                                                [1, 0, 0],
+                                                [1, 0, 0]], dtype=np.int32))
+    pestat = ops.convert_to_tensor(np.ones((3, 16), dtype=np.float32))
+    batch = ops.convert_to_tensor(np.array([0, 0, 0], dtype=np.int32))
+
+    pred_eval = model(x, edge_index, edge_attr=edge_attr, pestat_RWSE=pestat, batch=batch, training=False)
+    pred_train = model(x, edge_index, edge_attr=edge_attr, pestat_RWSE=pestat, batch=batch, training=True)
+
+    assert ops.shape(pred_eval) == (1, 1)
+    assert ops.shape(pred_train) == (1, 1)
+
+
+def test_gps_model_load_weights_synthetic():
+    torch = pytest.importorskip("torch")
+
+    dim_h = 32
+    rwse_pe = 8
+    model = GPSModel(
+        dim_in=dim_h,
+        dim_out=1,
+        num_layers=1,
+        dim_hidden=dim_h,
+        num_heads=2,
+        local_gnn_type="CustomGatedGCN",
+        act="gelu",
+        dropout=0.0,
+        attn_dropout=0.0,
+        batch_norm=True,
+        node_encoder_type="Atom+RWSE",
+        edge_encoder_type="Bond",
+        rwse_num_steps=16,
+        rwse_dim_pe=rwse_pe,
+        graph_pooling="mean",
+        head_layers=1,
+    )
+    model.build(None)
+
+    # Synthetic PyTorch state dict
+    state_dict = {}
+    for i in range(9):
+        dim_val = [119, 4, 12, 12, 10, 6, 6, 2, 2][i]
+        state_dict[f"encoder.node_encoder.encoder1.atom_embedding_list.{i}.weight"] = torch.randn(dim_val, dim_h - rwse_pe)
+    for i in range(3):
+        dim_val = [5, 6, 2][i]
+        state_dict[f"encoder.edge_encoder.bond_embedding_list.{i}.weight"] = torch.randn(dim_val, dim_h)
+
+    p_rwse = "encoder.node_encoder.encoder2"
+    state_dict[f"{p_rwse}.raw_norm.weight"] = torch.ones(16)
+    state_dict[f"{p_rwse}.raw_norm.bias"] = torch.zeros(16)
+    state_dict[f"{p_rwse}.raw_norm.running_mean"] = torch.zeros(16)
+    state_dict[f"{p_rwse}.raw_norm.running_var"] = torch.ones(16)
+    state_dict[f"{p_rwse}.pe_encoder.weight"] = torch.randn(rwse_pe, 16)
+    state_dict[f"{p_rwse}.pe_encoder.bias"] = torch.zeros(rwse_pe)
+
+    p0 = "layers.0"
+    for proj in ["A", "B", "C", "D", "E"]:
+        state_dict[f"{p0}.local_model.{proj}.weight"] = torch.randn(dim_h, dim_h)
+        state_dict[f"{p0}.local_model.{proj}.bias"] = torch.zeros(dim_h)
+    for bn in ["bn_node_x", "bn_edge_e"]:
+        state_dict[f"{p0}.local_model.{bn}.weight"] = torch.ones(dim_h)
+        state_dict[f"{p0}.local_model.{bn}.bias"] = torch.zeros(dim_h)
+        state_dict[f"{p0}.local_model.{bn}.running_mean"] = torch.zeros(dim_h)
+        state_dict[f"{p0}.local_model.{bn}.running_var"] = torch.ones(dim_h)
+
+    state_dict[f"{p0}.norm1_local.weight"] = torch.ones(dim_h)
+    state_dict[f"{p0}.norm1_local.bias"] = torch.zeros(dim_h)
+    state_dict[f"{p0}.norm1_local.running_mean"] = torch.zeros(dim_h)
+    state_dict[f"{p0}.norm1_local.running_var"] = torch.ones(dim_h)
+
+    state_dict[f"{p0}.self_attn.in_proj_weight"] = torch.randn(3 * dim_h, dim_h)
+    state_dict[f"{p0}.self_attn.in_proj_bias"] = torch.zeros(3 * dim_h)
+    state_dict[f"{p0}.self_attn.out_proj.weight"] = torch.randn(dim_h, dim_h)
+    state_dict[f"{p0}.self_attn.out_proj.bias"] = torch.zeros(dim_h)
+
+    state_dict[f"{p0}.norm1_attn.weight"] = torch.ones(dim_h)
+    state_dict[f"{p0}.norm1_attn.bias"] = torch.zeros(dim_h)
+    state_dict[f"{p0}.norm1_attn.running_mean"] = torch.zeros(dim_h)
+    state_dict[f"{p0}.norm1_attn.running_var"] = torch.ones(dim_h)
+
+    state_dict[f"{p0}.ff_linear1.weight"] = torch.randn(dim_h * 2, dim_h)
+    state_dict[f"{p0}.ff_linear1.bias"] = torch.zeros(dim_h * 2)
+    state_dict[f"{p0}.ff_linear2.weight"] = torch.randn(dim_h, dim_h * 2)
+    state_dict[f"{p0}.ff_linear2.bias"] = torch.zeros(dim_h)
+
+    state_dict[f"{p0}.norm2.weight"] = torch.ones(dim_h)
+    state_dict[f"{p0}.norm2.bias"] = torch.zeros(dim_h)
+    state_dict[f"{p0}.norm2.running_mean"] = torch.zeros(dim_h)
+    state_dict[f"{p0}.norm2.running_var"] = torch.ones(dim_h)
+
+    state_dict["post_mp.FC_layers.0.weight"] = torch.randn(dim_h // 2, dim_h)
+    state_dict["post_mp.FC_layers.0.bias"] = torch.zeros(dim_h // 2)
+    state_dict["post_mp.FC_layers.1.weight"] = torch.randn(1, dim_h // 2)
+    state_dict["post_mp.FC_layers.1.bias"] = torch.zeros(1)
+
+    with tempfile.NamedTemporaryFile(suffix=".ckpt", delete=False) as f:
+        torch.save({"model_state": state_dict}, f.name)
+        ckpt_path = f.name
+
+    try:
+        load_gps_weights(model, ckpt_path)
+    finally:
+        if os.path.exists(ckpt_path):
+            os.remove(ckpt_path)
+
+
+def test_gps_model_pretrained_checkpoint_if_available():
+    torch = pytest.importorskip("torch")
+    ckpt_path = download_gps_checkpoint("pcqm4m-GPS+RWSE.deep")
+    if os.path.exists(ckpt_path):
+        model = GPSModel(
+            dim_in=256,
+            dim_out=1,
+            num_layers=16,
+            dim_hidden=256,
+            num_heads=8,
+            local_gnn_type="CustomGatedGCN",
+            act="gelu",
+            dropout=0.0,
+            attn_dropout=0.0,
+            batch_norm=True,
+            layer_norm=False,
+            node_encoder_type="Atom+RWSE",
+            edge_encoder_type="Bond",
+            rwse_num_steps=16,
+            rwse_dim_pe=20,
+            graph_pooling="mean",
+            head_layers=2,
+        )
+        x = ops.convert_to_tensor(np.array([[6, 0, 4, 4, 3, 2, 2, 0, 0],
+                                            [8, 0, 2, 2, 2, 1, 1, 0, 0],
+                                            [6, 0, 3, 3, 3, 2, 2, 0, 0]], dtype=np.int32))
+        edge_index = ops.convert_to_tensor(np.array([[0, 1, 1, 2],
+                                                     [1, 0, 2, 1]], dtype=np.int32))
+        edge_attr = ops.convert_to_tensor(np.array([[0, 0, 0],
+                                                    [0, 0, 0],
+                                                    [1, 0, 0],
+                                                    [1, 0, 0]], dtype=np.int32))
+        pestat = ops.convert_to_tensor(np.ones((3, 16), dtype=np.float32))
+        batch = ops.convert_to_tensor(np.array([0, 0, 0], dtype=np.int32))
+
+        # Build and load
+        _ = model(x, edge_index, edge_attr=edge_attr, pestat_RWSE=pestat, batch=batch, training=False)
+        load_gps_weights(model, ckpt_path)
+
+        out = model(x, edge_index, edge_attr=edge_attr, pestat_RWSE=pestat, batch=batch, training=False)
+        val = float(ops.convert_to_numpy(out)[0, 0])
+        # PyTorch reference output was ~15.3017
+        np.testing.assert_allclose(val, 15.3017, rtol=1e-3, atol=1e-3)
+

k3_node/models/test_gpse.py

@@ -0,0 +1,34 @@
+import keras.ops as ops
+from k3_node.models.gpse import GPSE, GPSENodeEncoder
+
+
+def test_gpse_encoder():
+    x = ops.ones((4, 16))
+    pos_enc = ops.ones((4, 32))
+
+    encoder = GPSENodeEncoder(
+        dim_emb=64,
+        dim_pe_in=32,
+        dim_pe_out=16,
+        dim_in=16,
+        expand_x=True,
+    )
+    out = encoder(x, pos_enc)
+    assert out.shape == (4, 64)
+
+
+def test_gpse_model():
+    x = ops.ones((4, 16))
+    edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]])
+
+    model = GPSE(
+        dim_in=16,
+        dim_inner=32,
+        layers_pre_mp=1,
+        layers_mp=2,
+        layers_post_mp=1,
+        use_repr=True,
+    )
+    out = model(x, edge_index)
+    assert out.shape == (4, 32)
+