k3-node 1.0.0__py3-none-any.whl

k3_node/models/test_graph_unet.py

@@ -0,0 +1,26 @@
+from keras import ops, random
+from k3_node.models import GraphUNet
+
+
+def test_graph_unet():
+    model = GraphUNet(
+        in_channels=16,
+        hidden_channels=32,
+        out_channels=8,
+        depth=2,
+        pool_ratios=0.5,
+    )
+    x = random.normal((6, 16))
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 4, 5, 0, 2],
+        [1, 2, 0, 4, 5, 3, 3, 4],
+    ], dtype="int64")
+
+    out = model(x, edge_index)
+    assert out.shape == (6, 8)
+
+    # Test with batch
+    batch = ops.convert_to_tensor([0, 0, 0, 1, 1, 1], dtype="int64")
+    out_batch = model(x, edge_index, batch=batch)
+    assert out_batch.shape == (6, 8)
+

k3_node/models/test_graphmae2.py

@@ -0,0 +1,226 @@
+import os
+import pytest
+import numpy as np
+from keras import ops
+
+from k3_node.models import (
+    GraphMAE2,
+    sce_loss,
+    load_graphmae2_weights,
+    download_graphmae2_checkpoint,
+)
+
+
+def test_sce_loss():
+    # Identical vectors -> loss should be 0
+    x = ops.convert_to_tensor([[1.0, 0.0], [0.0, 1.0]], dtype="float32")
+    loss_identical = sce_loss(x, x, alpha=2.0)
+    assert float(ops.convert_to_numpy(loss_identical)) < 1e-6
+
+    # Orthogonal vectors -> cos_sim = 0, (1 - 0)^2 = 1
+    y = ops.convert_to_tensor([[0.0, 1.0], [1.0, 0.0]], dtype="float32")
+    loss_ortho = sce_loss(x, y, alpha=2.0)
+    assert abs(float(ops.convert_to_numpy(loss_ortho)) - 1.0) < 1e-5
+
+    # Opposing vectors -> cos_sim = -1, (1 - (-1))^2 = 4
+    z = ops.convert_to_tensor([[-1.0, 0.0], [0.0, -1.0]], dtype="float32")
+    loss_opp = sce_loss(x, z, alpha=2.0)
+    assert abs(float(ops.convert_to_numpy(loss_opp)) - 4.0) < 1e-5
+
+
+def test_graphmae2_basic():
+    in_dim = 16
+    num_hidden = 32
+    num_nodes = 8
+
+    x = ops.convert_to_tensor(np.random.randn(num_nodes, in_dim).astype("float32"))
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 4, 5, 6, 7, 0, 2],
+        [1, 2, 3, 4, 5, 6, 7, 0, 2, 0],
+    ], dtype="int32")
+
+    model = GraphMAE2(
+        in_dim=in_dim,
+        num_hidden=num_hidden,
+        num_layers=2,
+        num_dec_layers=1,
+        num_remasking=2,
+        nhead=4,
+        nhead_out=1,
+        activation="prelu",
+        norm="layernorm",
+        residual=True,
+    )
+
+    expected_repr = f"GraphMAE2(in_dim={in_dim}, num_hidden={num_hidden}, num_layers=2, num_dec_layers=1, nhead=4)"
+    assert str(model) == expected_repr
+
+    # Forward pass / embed
+    emb = model(x, edge_index)
+    assert ops.shape(emb) == (num_nodes, num_hidden)
+
+    emb2 = model.embed(x, edge_index)
+    assert ops.shape(emb2) == (num_nodes, num_hidden)
+
+    # Loss computation
+    loss_val = model.loss(x, edge_index)
+    assert float(ops.convert_to_numpy(loss_val)) >= 0.0
+
+    # EMA update
+    initial_teacher_w = float(ops.convert_to_numpy(ops.mean(model.encoder_ema.gat_layers[0].fc.kernel)))
+    # Modify student weights
+    model.encoder.gat_layers[0].fc.kernel.assign(model.encoder.gat_layers[0].fc.kernel + 1.0)
+    model.ema_update(momentum=0.5)
+    updated_teacher_w = float(ops.convert_to_numpy(ops.mean(model.encoder_ema.gat_layers[0].fc.kernel)))
+    assert updated_teacher_w != initial_teacher_w
+
+
+@pytest.mark.parametrize("remask_method", ["random", "fixed"])
+def test_graphmae2_remask_methods(remask_method):
+    in_dim = 8
+    num_hidden = 16
+    num_nodes = 6
+
+    x = ops.convert_to_tensor(np.random.randn(num_nodes, in_dim).astype("float32"))
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 4, 5],
+        [1, 2, 3, 4, 5, 0],
+    ], dtype="int32")
+
+    model = GraphMAE2(
+        in_dim=in_dim,
+        num_hidden=num_hidden,
+        num_layers=2,
+        num_dec_layers=1,
+        num_remasking=2,
+        nhead=2,
+        nhead_out=1,
+        remask_method=remask_method,
+        loss_fn="sce",
+    )
+
+    loss_val = model.loss(x, edge_index)
+    assert float(ops.convert_to_numpy(loss_val)) >= 0.0
+
+
+def test_load_graphmae2_checkpoint():
+    ckpt_path = "GraphMAE2-main/GraphMAE2_checkpoints/gat_gat_1024_4_ogbn-arxiv_0.5_1024_checkpoint.pt"
+    if not os.path.exists(ckpt_path):
+        pytest.skip(f"Checkpoint {ckpt_path} not found")
+
+    model = GraphMAE2(
+        in_dim=128,
+        num_hidden=1024,
+        num_layers=4,
+        num_dec_layers=1,
+        nhead=8,
+        nhead_out=1,
+        activation="prelu",
+        norm="layernorm",
+        residual=True,
+    )
+
+    load_graphmae2_weights(model, ckpt_path)
+
+    # Verify forward pass with loaded weights
+    num_nodes = 5
+    x = ops.convert_to_tensor(np.random.randn(num_nodes, 128).astype("float32"))
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 4, 0],
+        [1, 2, 3, 4, 0, 2],
+    ], dtype="int32")
+
+    out = model(x, edge_index)
+    assert ops.shape(out) == (num_nodes, 1024)
+
+    # Compute loss with loaded weights
+    loss = model.loss(x, edge_index)
+    assert float(ops.convert_to_numpy(loss)) >= 0.0
+
+
+def test_load_graphmae2_checkpoint_products():
+    ckpt_path = "GraphMAE2-main/GraphMAE2_checkpoints/gat_gat_1024_4_ogbn-products_0.5_1024_checkpoint.pt"
+    if not os.path.exists(ckpt_path):
+        pytest.skip(f"Checkpoint {ckpt_path} not found")
+
+    model = GraphMAE2(
+        in_dim=100,
+        num_hidden=1024,
+        num_layers=4,
+        num_dec_layers=1,
+        nhead=4,
+        nhead_out=1,
+        activation="prelu",
+        norm="layernorm",
+        residual=True,
+    )
+
+    model.load_weights_from_checkpoint(ckpt_path)
+
+    num_nodes = 4
+    x = ops.convert_to_tensor(np.random.randn(num_nodes, 100).astype("float32"))
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3],
+        [1, 2, 3, 0],
+    ], dtype="int32")
+
+    out = model(x, edge_index)
+    assert ops.shape(out) == (num_nodes, 1024)
+
+
+def test_download_graphmae2_checkpoint(monkeypatch, tmp_path):
+    # Test invalid dataset
+    with pytest.raises(ValueError, match="Unknown dataset"):
+        download_graphmae2_checkpoint("unknown_dataset")
+
+    # Test existing file returns path directly without download
+    local_dir = "GraphMAE2-main/GraphMAE2_checkpoints"
+    if os.path.exists(local_dir):
+        path = download_graphmae2_checkpoint("ogbn-arxiv", folder=local_dir)
+        assert os.path.exists(path)
+        assert "ogbn-arxiv" in path
+
+    # Test download invocation using monkeypatched download_google_url
+    called = {}
+
+    def mock_download(id, folder, filename, log=True):
+        called["id"] = id
+        called["folder"] = folder
+        called["filename"] = filename
+        os.makedirs(folder, exist_ok=True)
+        target = os.path.join(folder, filename)
+        with open(target, "wb") as f:
+            f.write(b"mock_content")
+        return target
+
+    monkeypatch.setattr("k3_node.models.graphmae2.download_google_url", mock_download)
+
+    fake_folder = str(tmp_path / "ckpts")
+    res_path = download_graphmae2_checkpoint("mag-scholar-f", folder=fake_folder)
+    assert os.path.exists(res_path)
+    assert called["id"] == "1KpQk_OKbbo4qTLQYZ84pAJDy1sh4oZv2"
+    assert "mag-scholar-f" in called["filename"]
+
+
+def test_graphmae2_from_pretrained():
+    local_dir = "GraphMAE2-main/GraphMAE2_checkpoints"
+    if not os.path.exists(local_dir):
+        pytest.skip(f"Checkpoints directory {local_dir} not found")
+
+    model = GraphMAE2.from_pretrained("ogbn-arxiv", folder=local_dir)
+    assert model.in_dim == 128
+    assert model.num_hidden == 1024
+    assert model.nhead == 8
+
+    num_nodes = 3
+    x = ops.convert_to_tensor(np.random.randn(num_nodes, 128).astype("float32"))
+    edge_index = ops.convert_to_tensor([[0, 1, 2], [1, 2, 0]], dtype="int32")
+    out = model(x, edge_index)
+    assert ops.shape(out) == (num_nodes, 1024)
+
+    # Test products config
+    model_prod = GraphMAE2.from_pretrained("ogbn-products", folder=local_dir)
+    assert model_prod.in_dim == 100
+    assert model_prod.num_hidden == 1024
+    assert model_prod.nhead == 4
+

k3_node/models/test_graphormer.py

@@ -0,0 +1,233 @@
+import os
+import pytest
+import numpy as np
+try:
+    import torch
+except ImportError:
+    torch = None
+from keras import ops
+
+from k3_node.models import (
+    Graphormer,
+    GraphNodeFeature,
+    GraphAttnBias,
+    GraphormerMultiheadAttention,
+    GraphormerGraphEncoderLayer,
+    GraphormerGraphEncoder,
+    load_graphormer_weights,
+    download_graphormer_checkpoint,
+)
+from k3_node.models.graphormer import get_graphormer_config
+
+
+def test_graph_node_feature():
+    num_atoms, num_in_deg, num_out_deg, hidden_dim = 16, 10, 10, 32
+    gnf = GraphNodeFeature(
+        num_atoms=num_atoms,
+        num_in_degree=num_in_deg,
+        num_out_degree=num_out_deg,
+        hidden_dim=hidden_dim,
+    )
+
+    batch_size, num_nodes = 2, 5
+    x = ops.convert_to_tensor(np.random.randint(1, num_atoms, size=(batch_size, num_nodes)), dtype="int32")
+    in_deg = ops.convert_to_tensor(np.random.randint(0, num_in_deg, size=(batch_size, num_nodes)), dtype="int32")
+    out_deg = ops.convert_to_tensor(np.random.randint(0, num_out_deg, size=(batch_size, num_nodes)), dtype="int32")
+
+    out = gnf(x, in_deg, out_deg)
+    # Shape should be [batch_size, num_nodes + 1, hidden_dim]
+    assert ops.shape(out) == (batch_size, num_nodes + 1, hidden_dim)
+
+    # Test multi-dimensional atom features [B, N, D]
+    x_multi = ops.convert_to_tensor(np.random.randint(1, num_atoms, size=(batch_size, num_nodes, 3)), dtype="int32")
+    out_multi = gnf(x_multi, in_deg, out_deg)
+    assert ops.shape(out_multi) == (batch_size, num_nodes + 1, hidden_dim)
+
+
+def test_graph_attn_bias():
+    num_heads, num_atoms, num_edges, num_spatial, num_edge_dis = 4, 16, 8, 10, 5
+    gab = GraphAttnBias(
+        num_heads=num_heads,
+        num_atoms=num_atoms,
+        num_edges=num_edges,
+        num_spatial=num_spatial,
+        num_edge_dis=num_edge_dis,
+        edge_type="multi_hop",
+    )
+
+    batch_size, num_nodes = 2, 4
+    attn_bias = ops.zeros((batch_size, num_nodes + 1, num_nodes + 1), dtype="float32")
+    spatial_pos = ops.convert_to_tensor(
+        np.random.randint(0, num_spatial, size=(batch_size, num_nodes, num_nodes)), dtype="int32"
+    )
+    x = ops.zeros((batch_size, num_nodes, 1), dtype="int32")
+    edge_input = ops.convert_to_tensor(
+        np.random.randint(0, num_edges, size=(batch_size, num_nodes, num_nodes, 3, 2)), dtype="int32"
+    )
+
+    out = gab(attn_bias=attn_bias, spatial_pos=spatial_pos, x=x, edge_input=edge_input)
+    assert ops.shape(out) == (batch_size, num_heads, num_nodes + 1, num_nodes + 1)
+
+    # Test single-hop edge bias
+    gab_single = GraphAttnBias(
+        num_heads=num_heads,
+        num_atoms=num_atoms,
+        num_edges=num_edges,
+        num_spatial=num_spatial,
+        num_edge_dis=num_edge_dis,
+        edge_type="single_hop",
+    )
+    attn_edge_type = ops.convert_to_tensor(
+        np.random.randint(0, num_edges, size=(batch_size, num_nodes, num_nodes, 2)), dtype="int32"
+    )
+    out_single = gab_single(
+        attn_bias=attn_bias, spatial_pos=spatial_pos, x=x, attn_edge_type=attn_edge_type
+    )
+    assert ops.shape(out_single) == (batch_size, num_heads, num_nodes + 1, num_nodes + 1)
+
+
+def test_graphormer_attention_and_layer():
+    embed_dim, num_heads = 32, 4
+    mha = GraphormerMultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=0.0)
+
+    batch_size, seq_len = 2, 6
+    x = ops.convert_to_tensor(np.random.randn(batch_size, seq_len, embed_dim).astype("float32"))
+    attn_bias = ops.zeros((batch_size, num_heads, seq_len, seq_len), dtype="float32")
+    key_padding_mask = ops.convert_to_tensor([[False, False, False, False, False, True],
+                                              [False, False, False, False, True, True]], dtype="bool")
+
+    out, weights = mha(x, attn_bias=attn_bias, key_padding_mask=key_padding_mask)
+    assert ops.shape(out) == (batch_size, seq_len, embed_dim)
+    assert ops.shape(weights) == (batch_size, num_heads, seq_len, seq_len)
+
+    # Test encoder layers (both Pre-LN and Post-LN)
+    for pre_ln in [True, False]:
+        enc_layer = GraphormerGraphEncoderLayer(
+            embedding_dim=embed_dim,
+            ffn_embedding_dim=64,
+            num_attention_heads=num_heads,
+            pre_layernorm=pre_ln,
+        )
+        layer_out = enc_layer(x, attn_bias=attn_bias, key_padding_mask=key_padding_mask)
+        assert ops.shape(layer_out) == (batch_size, seq_len, embed_dim)
+
+
+def test_graphormer_model():
+    batch_size, num_nodes = 2, 4
+    model = Graphormer(
+        num_atoms=16,
+        num_in_degree=8,
+        num_out_degree=8,
+        num_edges=8,
+        num_spatial=8,
+        num_edge_dis=4,
+        num_encoder_layers=2,
+        embedding_dim=32,
+        ffn_embedding_dim=64,
+        num_attention_heads=4,
+        num_classes=1,
+    )
+
+    data = {
+        "x": ops.convert_to_tensor(np.random.randint(1, 15, size=(batch_size, num_nodes, 2)), dtype="int32"),
+        "in_degree": ops.convert_to_tensor(np.random.randint(0, 7, size=(batch_size, num_nodes)), dtype="int32"),
+        "out_degree": ops.convert_to_tensor(np.random.randint(0, 7, size=(batch_size, num_nodes)), dtype="int32"),
+        "attn_bias": ops.zeros((batch_size, num_nodes + 1, num_nodes + 1), dtype="float32"),
+        "spatial_pos": ops.convert_to_tensor(np.random.randint(0, 7, size=(batch_size, num_nodes, num_nodes)), dtype="int32"),
+        "edge_input": ops.convert_to_tensor(np.random.randint(0, 7, size=(batch_size, num_nodes, num_nodes, 2, 2)), dtype="int32"),
+    }
+
+    # Forward via batched_data dict
+    out = model(data)
+    assert ops.shape(out) == (batch_size, 1)
+
+    # Forward with return_all=True
+    pred, all_h = model(data, return_all=True)
+    assert ops.shape(pred) == (batch_size, 1)
+    assert ops.shape(all_h) == (batch_size, num_nodes + 1, 32)
+
+    # Embed method
+    emb = model.embed(data)
+    assert ops.shape(emb) == (batch_size, 32)
+
+    # Keyword argument forward
+    out_kwargs = model(
+        x=data["x"],
+        in_degree=data["in_degree"],
+        out_degree=data["out_degree"],
+        attn_bias=data["attn_bias"],
+        spatial_pos=data["spatial_pos"],
+        edge_input=data["edge_input"],
+    )
+    assert ops.shape(out_kwargs) == (batch_size, 1)
+
+
+def test_graphormer_presets():
+    for variant in ["base", "slim", "large"]:
+        cfg = get_graphormer_config(variant)
+        assert "num_encoder_layers" in cfg
+        assert "embedding_dim" in cfg
+        assert "num_attention_heads" in cfg
+
+
+def test_download_and_load_checkpoint(monkeypatch, tmp_path):
+    with pytest.raises(ValueError, match="Unknown pretrained model"):
+        download_graphormer_checkpoint("non_existent_model")
+
+    called = {}
+
+    def mock_download(url, folder, filename, log=True):
+        called["url"] = url
+        called["folder"] = folder
+        called["filename"] = filename
+        os.makedirs(folder, exist_ok=True)
+        target = os.path.join(folder, filename)
+        with open(target, "wb") as f:
+            f.write(b"mock")
+        return target
+
+    monkeypatch.setattr("k3_node.data.download.download_url", mock_download)
+
+    ckpt_folder = str(tmp_path / "ckpts")
+    path = download_graphormer_checkpoint("pcqm4mv1_graphormer_base", folder=ckpt_folder)
+    assert os.path.exists(path)
+    assert "pcqm4mv1" in called["filename"]
+
+    # Test loading synthetic checkpoint
+    if torch is None:
+        return
+
+    model = Graphormer(
+        num_atoms=8,
+        num_in_degree=4,
+        num_out_degree=4,
+        num_edges=4,
+        num_spatial=4,
+        num_edge_dis=2,
+        num_encoder_layers=1,
+        embedding_dim=16,
+        ffn_embedding_dim=32,
+        num_attention_heads=2,
+        num_classes=1,
+    )
+    model.build(None)
+
+    fake_state = {
+        "encoder.graph_encoder.graph_node_feature.atom_encoder.weight": torch.randn(9, 16),
+        "encoder.graph_encoder.graph_node_feature.in_degree_encoder.weight": torch.randn(4, 16),
+        "encoder.graph_encoder.graph_node_feature.out_degree_encoder.weight": torch.randn(4, 16),
+        "encoder.graph_encoder.graph_node_feature.graph_token.weight": torch.randn(1, 16),
+        "encoder.graph_encoder.layers.0.fc1.weight": torch.randn(32, 16),
+        "encoder.graph_encoder.layers.0.fc1.bias": torch.randn(32),
+        "encoder.graph_encoder.layers.0.fc2.weight": torch.randn(16, 32),
+        "encoder.graph_encoder.layers.0.fc2.bias": torch.randn(16),
+        "encoder.lm_head_transform_weight.weight": torch.randn(16, 16),
+        "encoder.lm_head_transform_weight.bias": torch.randn(16),
+        "encoder.embed_out.weight": torch.randn(1, 16),
+        "encoder.lm_output_learned_bias": torch.tensor([0.42]),
+    }
+    ckpt_file = str(tmp_path / "fake_ckpt.pt")
+    torch.save({"model": fake_state}, ckpt_file)
+
+    loaded_model = load_graphormer_weights(model, checkpoint_path=ckpt_file)
+    assert float(ops.convert_to_numpy(loaded_model.lm_output_learned_bias[0])) == pytest.approx(0.42, abs=1e-5)

k3_node/models/test_graphormer3d.py

@@ -0,0 +1,163 @@
+import os
+import pytest
+import numpy as np
+try:
+    import torch
+except ImportError:
+    torch = None
+from keras import ops
+
+from k3_node.models import (
+    Graphormer3D,
+    GaussianLayer,
+    RBF,
+    Graphormer3DEncoderLayer,
+    NodeTaskHead,
+    load_graphormer3d_weights,
+    download_graphormer3d_checkpoint,
+)
+
+
+def test_gaussian_layer():
+    num_kernel, edge_types = 32, 16
+    gbf = GaussianLayer(num_kernel=num_kernel, edge_types=edge_types)
+
+    batch_size, num_nodes = 2, 4
+    dist = ops.convert_to_tensor(np.abs(np.random.randn(batch_size, num_nodes, num_nodes)).astype("float32"))
+    edges = ops.convert_to_tensor(np.random.randint(0, edge_types, size=(batch_size, num_nodes, num_nodes)), dtype="int32")
+
+    out = gbf(dist, edges)
+    assert ops.shape(out) == (batch_size, num_nodes, num_nodes, num_kernel)
+    assert np.all(ops.convert_to_numpy(out) >= 0.0)
+
+
+def test_rbf_layer():
+    num_kernel, edge_types = 32, 16
+    rbf = RBF(num_kernel=num_kernel, edge_types=edge_types)
+
+    batch_size, num_nodes = 2, 4
+    dist = ops.convert_to_tensor(np.abs(np.random.randn(batch_size, num_nodes, num_nodes)).astype("float32"))
+    edges = ops.convert_to_tensor(np.random.randint(0, edge_types, size=(batch_size, num_nodes, num_nodes)), dtype="int32")
+
+    out = rbf(dist, edges)
+    assert ops.shape(out) == (batch_size, num_nodes, num_nodes, num_kernel)
+    assert np.all(ops.convert_to_numpy(out) >= 0.0)
+
+
+def test_node_task_head():
+    embed_dim, num_heads = 32, 4
+    head = NodeTaskHead(embed_dim=embed_dim, num_heads=num_heads)
+
+    batch_size, num_nodes = 2, 5
+    query = ops.convert_to_tensor(np.random.randn(batch_size, num_nodes, embed_dim).astype("float32"))
+    attn_bias = ops.zeros((batch_size * num_heads, num_nodes, num_nodes), dtype="float32")
+    delta_pos = ops.convert_to_tensor(np.random.randn(batch_size, num_nodes, num_nodes, 3).astype("float32"))
+
+    forces = head(query, attn_bias, delta_pos)
+    assert ops.shape(forces) == (batch_size, num_nodes, 3)
+
+
+def test_graphormer3d_encoder_layer():
+    embed_dim, num_heads = 32, 4
+    layer = Graphormer3DEncoderLayer(
+        embedding_dim=embed_dim,
+        ffn_embedding_dim=64,
+        num_attention_heads=num_heads,
+    )
+
+    batch_size, num_nodes = 2, 4
+    x = ops.convert_to_tensor(np.random.randn(batch_size, num_nodes, embed_dim).astype("float32"))
+    attn_bias = ops.zeros((batch_size * num_heads, num_nodes, num_nodes), dtype="float32")
+
+    out = layer(x, attn_bias=attn_bias)
+    assert ops.shape(out) == (batch_size, num_nodes, embed_dim)
+
+
+def test_graphormer3d_model():
+    batch_size, num_nodes = 2, 5
+    model = Graphormer3D(
+        layers=2,
+        blocks=2,
+        embed_dim=32,
+        ffn_embed_dim=64,
+        attention_heads=4,
+        num_kernel=16,
+        atom_types=16,
+    )
+
+    atoms = ops.convert_to_tensor([[1, 2, 3, 4, 0], [2, 3, 4, 0, 0]], dtype="int32")
+    tags = ops.convert_to_tensor([[1, 1, 2, 2, 0], [1, 2, 2, 0, 0]], dtype="int32")
+    pos = ops.convert_to_tensor(np.random.randn(batch_size, num_nodes, 3).astype("float32"))
+
+    energy, forces = model(atoms, tags, pos)
+    assert ops.shape(energy) == (batch_size,)
+    assert ops.shape(forces) == (batch_size, num_nodes, 3)
+
+
+def test_graphormer3d_checkpoint_download_and_load(monkeypatch, tmp_path):
+    with pytest.raises(ValueError, match="Unknown pretrained 3D model"):
+        download_graphormer3d_checkpoint("invalid_3d_name")
+
+    called = {}
+
+    def mock_download(url, folder, filename, log=True):
+        called["url"] = url
+        called["folder"] = folder
+        called["filename"] = filename
+        os.makedirs(folder, exist_ok=True)
+        target = os.path.join(folder, filename)
+        with open(target, "wb") as f:
+            f.write(b"mock_3d")
+        return target
+
+    monkeypatch.setattr("k3_node.data.download.download_url", mock_download)
+
+    ckpt_folder = str(tmp_path / "ckpts3d")
+    path = download_graphormer3d_checkpoint("oc20is2re_graphormer3d_base", folder=ckpt_folder)
+    assert os.path.exists(path)
+    assert "oc20is2re" in called["filename"]
+
+    # Test loading synthetic checkpoint
+    if torch is None:
+        return
+
+    model = Graphormer3D(
+        layers=1,
+        blocks=1,
+        embed_dim=16,
+        ffn_embed_dim=32,
+        attention_heads=2,
+        num_kernel=8,
+        atom_types=8,
+    )
+    model.build(None)
+
+    fake_state = {
+        "encoder.atom_encoder.weight": torch.randn(8, 16),
+        "encoder.tag_encoder.weight": torch.randn(3, 16),
+        "encoder.gbf.means.weight": torch.randn(1, 8),
+        "encoder.gbf.stds.weight": torch.randn(1, 8),
+        "encoder.gbf.mul.weight": torch.randn(64, 1),
+        "encoder.gbf.bias.weight": torch.randn(64, 1),
+        "encoder.bias_proj.layer1.weight": torch.randn(8, 8),
+        "encoder.bias_proj.layer1.bias": torch.randn(8),
+        "encoder.bias_proj.layer2.weight": torch.randn(2, 8),
+        "encoder.bias_proj.layer2.bias": torch.randn(2),
+        "encoder.edge_proj.weight": torch.randn(16, 8),
+        "encoder.edge_proj.bias": torch.randn(16),
+        "encoder.layers.0.self_attn.in_proj.weight": torch.randn(48, 16),
+        "encoder.layers.0.self_attn.in_proj.bias": torch.randn(48),
+        "encoder.layers.0.self_attn.out_proj.weight": torch.randn(16, 16),
+        "encoder.layers.0.self_attn.out_proj.bias": torch.randn(16),
+        "encoder.engergy_proj.layer1.weight": torch.randn(16, 16),
+        "encoder.engergy_proj.layer1.bias": torch.randn(16),
+        "encoder.engergy_proj.layer2.weight": torch.randn(1, 16),
+        "encoder.engergy_proj.layer2.bias": torch.randn(1),
+        "encoder.energe_agg_factor.weight": torch.randn(3, 1),
+    }
+    ckpt_file = str(tmp_path / "fake_3d.pt")
+    torch.save({"model": fake_state}, ckpt_file)
+
+    loaded_model = load_graphormer3d_weights(model, checkpoint_path=ckpt_file)
+    assert loaded_model.atom_encoder.embeddings.shape == (8, 16)
+