k3-node 1.0.0__py3-none-any.whl

k3_node/models/test_rect.py

@@ -0,0 +1,23 @@
+from keras import ops
+
+from k3_node.models import RECT_L
+
+
+def test_rect():
+    x = ops.ones((6, 8))
+    y = 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 = RECT_L(8, 16)
+    assert str(model) == 'RECT_L(8, 16)'
+
+    out = model(x, edge_index)
+    assert ops.shape(out) == (6, 8)
+
+    embed_out = model.embed(x, edge_index)
+    assert ops.shape(embed_out) == (6, 16)
+
+    labels_out = model.get_semantic_labels(x, y, mask)
+    assert ops.shape(labels_out) == (3, 8)
+

k3_node/models/test_renet.py

@@ -0,0 +1,32 @@
+from keras import ops, random
+from k3_node.models import RENet
+
+
+def test_renet():
+    model = RENet(num_nodes=5, num_rels=4, hidden_channels=16, seq_len=3)
+
+    sub = ops.convert_to_tensor([0, 1], dtype="int64")
+    rel = ops.convert_to_tensor([0, 1], dtype="int64")
+    obj = ops.convert_to_tensor([2, 3], dtype="int64")
+
+    h_sub = ops.convert_to_tensor([0, 1, 2], dtype="int64")
+    h_sub_t = ops.convert_to_tensor([0, 1, 0], dtype="int64")
+    h_sub_batch = ops.convert_to_tensor([0, 0, 1], dtype="int64")
+
+    h_obj = ops.convert_to_tensor([1, 2, 3], dtype="int64")
+    h_obj_t = ops.convert_to_tensor([1, 2, 0], dtype="int64")
+    h_obj_batch = ops.convert_to_tensor([0, 0, 1], dtype="int64")
+
+    log_prob_obj, log_prob_sub = model(
+        sub, rel, obj,
+        h_sub, h_sub_t, h_sub_batch,
+        h_obj, h_obj_t, h_obj_batch,
+    )
+    assert log_prob_obj.shape == (2, 5)
+    assert log_prob_sub.shape == (2, 5)
+
+    # Test metrics
+    y = ops.convert_to_tensor([2, 3], dtype="int64")
+    metrics = model.test(log_prob_obj, y)
+    assert metrics.shape == (4,)
+

k3_node/models/test_rev_gnn.py

@@ -0,0 +1,24 @@
+import numpy as np
+from keras import ops, random
+from k3_node.layers.conv import GCNConv
+from k3_node.models import GroupAddRev
+
+
+def test_group_add_rev():
+    conv1 = GCNConv(16, 16)
+    conv2 = GCNConv(16, 16)
+    model = GroupAddRev([conv1, conv2])
+
+    x = random.normal((6, 32))
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 4],
+        [1, 2, 0, 4, 5],
+    ], dtype="int64")
+
+    out = model(x, edge_index=edge_index)
+    assert out.shape == (6, 32)
+
+    inv = model.inverse(out, edge_index=edge_index)
+    assert inv.shape == (6, 32)
+    assert np.allclose(ops.convert_to_numpy(x), ops.convert_to_numpy(inv), atol=1e-5)
+

k3_node/models/test_schnet.py

@@ -0,0 +1,43 @@
+import keras.ops as ops
+from k3_node.models.schnet import SchNet, ShiftedSoftplus, GaussianSmearing
+
+
+def test_shifted_softplus():
+    act = ShiftedSoftplus()
+    x = ops.convert_to_tensor([0.0, 1.0, 2.0])
+    out = act(x)
+    assert out.shape == (3,)
+    assert float(out[0]) == 0.0
+
+
+def test_gaussian_smearing():
+    smear = GaussianSmearing(start=0.0, stop=5.0, num_gaussians=10)
+    dist = ops.convert_to_tensor([0.0, 1.0, 2.5])
+    out = smear(dist)
+    assert out.shape == (3, 10)
+
+
+def test_schnet():
+    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 = SchNet(
+        hidden_channels=16,
+        num_filters=16,
+        num_interactions=2,
+        num_gaussians=10,
+        cutoff=5.0,
+    )
+    out = model(z, pos)
+    assert out.shape == (1, 1)
+
+    # Test with batch vector
+    batch = ops.convert_to_tensor([0, 0, 1, 1])
+    out_batched = model(z, pos, batch=batch)
+    assert out_batched.shape == (2, 1)
+

k3_node/models/test_sgformer.py

@@ -0,0 +1,48 @@
+from keras import ops, random
+from k3_node.layers.attention import SGFormerAttention
+from k3_node.models import SGFormer
+
+
+def test_sgformer_attention():
+    attn = SGFormerAttention(channels=32, heads=2, head_channels=16)
+    x = random.normal((2, 4, 32))
+    mask = ops.convert_to_tensor([[True, True, True, False], [True, True, False, False]])
+
+    out = attn(x, mask=mask)
+    assert out.shape == (2, 4, 16)
+
+
+def test_sgformer():
+    model = SGFormer(
+        in_channels=16,
+        hidden_channels=32,
+        out_channels=7,
+        trans_num_layers=1,
+        trans_num_heads=1,
+        gnn_num_layers=1,
+        aggregate="add",
+    )
+    x = random.normal((6, 16))
+    edge_index = ops.convert_to_tensor([[0, 1, 2, 3, 4], [1, 2, 0, 4, 5]], dtype="int64")
+
+    out = model(x, edge_index)
+    assert out.shape == (6, 7)
+
+    # 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, 7)
+
+    # Test aggregate='cat'
+    model_cat = SGFormer(
+        in_channels=16,
+        hidden_channels=32,
+        out_channels=7,
+        trans_num_layers=1,
+        trans_num_heads=1,
+        gnn_num_layers=1,
+        aggregate="cat",
+    )
+    out_cat = model_cat(x, edge_index)
+    assert out_cat.shape == (6, 7)
+

k3_node/models/test_signed_gcn.py

@@ -0,0 +1,28 @@
+import keras
+from keras import ops, random
+from k3_node.models import SignedGCN
+
+
+def test_signed_gcn():
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 0, 1, 4, 5],
+        [1, 2, 3, 0, 4, 5, 2, 3],
+    ], dtype="int64")
+
+    pos_edge_index, neg_edge_index = SignedGCN.split_edges(edge_index, test_ratio=0.2)
+    assert pos_edge_index.shape[0] == 2
+    assert neg_edge_index.shape[0] == 2
+
+    model = SignedGCN(in_channels=16, hidden_channels=32, num_layers=2)
+    x = random.normal((6, 16))
+
+    out = model(x, pos_edge_index, neg_edge_index)
+    assert out.shape == (6, 32)
+
+    # Test discriminate
+    disc = model.discriminate(out, edge_index)
+    assert disc.shape == (8, 3)
+
+    # Test losses
+    loss = model.loss(out, pos_edge_index, neg_edge_index)
+    assert loss.shape == ()

k3_node/models/test_tgn.py

@@ -0,0 +1,77 @@
+from keras import ops, random
+from k3_node.models import (
+    TGNMemory,
+    IdentityMessage,
+    LastAggregator,
+    MeanAggregator,
+    TimeEncoder,
+    LastNeighborLoader,
+)
+
+
+def test_time_encoder():
+    enc = TimeEncoder(out_channels=16)
+    t = ops.convert_to_tensor([1.0, 2.0, 3.0], dtype="float32")
+    out = enc(t)
+    assert out.shape == (3, 16)
+
+
+def test_aggregators():
+    msg = ops.convert_to_tensor([[1.0, 2.0], [3.0, 4.0]], dtype="float32")
+    index = ops.convert_to_tensor([0, 0], dtype="int64")
+    t = ops.convert_to_tensor([1.0, 2.0], dtype="float32")
+
+    last_aggr = LastAggregator()
+    out_last = last_aggr(msg, index, t, dim_size=1)
+    assert out_last.shape == (1, 2)
+    assert float(out_last[0, 0]) == 3.0
+
+    mean_aggr = MeanAggregator()
+    out_mean = mean_aggr(msg, index, t, dim_size=1)
+    assert out_mean.shape == (1, 2)
+    assert float(out_mean[0, 0]) == 2.0
+
+
+def test_last_neighbor_loader():
+    loader = LastNeighborLoader(num_nodes=4, size=2)
+    src = ops.convert_to_tensor([0, 1], dtype="int64")
+    dst = ops.convert_to_tensor([1, 2], dtype="int64")
+    loader.insert(src, dst)
+
+    n_id = ops.convert_to_tensor([0, 1], dtype="int64")
+    unique_nodes, edge_index, e_ids = loader(n_id)
+    assert len(unique_nodes.shape) == 1
+    assert edge_index.shape[0] == 2
+
+
+def test_tgn_memory():
+    memory_dim = 16
+    raw_msg_dim = 8
+    time_dim = 16
+    msg_module = IdentityMessage(raw_msg_dim, memory_dim, time_dim)
+    aggr_module = LastAggregator()
+
+    memory = TGNMemory(
+        num_nodes=5,
+        raw_msg_dim=raw_msg_dim,
+        memory_dim=memory_dim,
+        time_dim=time_dim,
+        message_module=msg_module,
+        aggregator_module=aggr_module,
+    )
+
+    n_id = ops.convert_to_tensor([0, 1, 2], dtype="int64")
+    mem, last_update = memory(n_id)
+    assert mem.shape == (3, memory_dim)
+    assert last_update.shape == (3,)
+
+    # Update state
+    src = ops.convert_to_tensor([0, 1], dtype="int64")
+    dst = ops.convert_to_tensor([1, 2], dtype="int64")
+    t = ops.convert_to_tensor([1.0, 2.0], dtype="float32")
+    raw_msg = random.normal((2, raw_msg_dim))
+
+    memory.update_state(src, dst, t, raw_msg)
+    mem_updated, last_updated = memory(n_id)
+    assert mem_updated.shape == (3, memory_dim)
+

k3_node/models/test_unimol.py

@@ -0,0 +1,179 @@
+import os
+import tempfile
+import numpy as np
+import pytest
+import torch
+import keras
+from keras import ops
+
+from k3_node.models import (
+    UniMolModel,
+    UniMolConfGenModel,
+    UniMolDockingModel,
+    UniMolGaussianLayer,
+    UniMolNonLinearHead,
+    UniMolDistanceHead,
+    download_unimol_checkpoint,
+    load_unimol_weights,
+)
+
+
+def test_unimol_gaussian_layer():
+    num_kernel = 64
+    edge_types = 128
+    layer = UniMolGaussianLayer(num_kernel=num_kernel, edge_types=edge_types)
+
+    bsz, seq_len = 2, 5
+    dist = ops.convert_to_tensor(np.random.uniform(0.5, 5.0, (bsz, seq_len, seq_len)).astype("float32"))
+    edge_type = ops.convert_to_tensor(np.random.randint(0, edge_types, (bsz, seq_len, seq_len)), dtype="int32")
+
+    out = layer(dist, edge_type)
+    assert ops.shape(out) == (bsz, seq_len, seq_len, num_kernel)
+
+
+def test_unimol_nonlinear_head():
+    head = UniMolNonLinearHead(input_dim=32, out_dim=16, activation_fn="gelu")
+    x = ops.convert_to_tensor(np.random.randn(2, 5, 32).astype("float32"))
+    out = head(x)
+    assert ops.shape(out) == (2, 5, 16)
+
+
+def test_unimol_distance_head():
+    head = UniMolDistanceHead(heads=8, activation_fn="gelu")
+    pair = ops.convert_to_tensor(np.random.randn(2, 6, 6, 8).astype("float32"))
+    dist = head(pair)
+    assert ops.shape(dist) == (2, 6, 6)
+    # Check symmetry
+    dist_np = ops.convert_to_numpy(dist)
+    np.testing.assert_allclose(dist_np, np.transpose(dist_np, (0, 2, 1)), atol=1e-5)
+
+
+def test_unimol_model_forward():
+    bsz = 2
+    seq_len = 8
+    vocab_size = 64
+    embed_dim = 64
+    heads = 4
+    layers_num = 2
+
+    model = UniMolModel(
+        output_dim=2,
+        vocab_size=vocab_size,
+        encoder_layers=layers_num,
+        encoder_embed_dim=embed_dim,
+        encoder_ffn_embed_dim=128,
+        encoder_attention_heads=heads,
+        num_kernel=32,
+    )
+
+    src_tokens = ops.convert_to_tensor(np.random.randint(1, vocab_size, (bsz, seq_len)), dtype="int32")
+    src_coord = ops.convert_to_tensor(np.random.randn(bsz, seq_len, 3).astype("float32"))
+
+    # Forward logits
+    logits = model(src_tokens, src_coord=src_coord)
+    assert ops.shape(logits) == (bsz, 2)
+
+    # Return repr
+    reprs = model(src_tokens, src_coord=src_coord, return_repr=True)
+    assert "cls_repr" in reprs
+    assert ops.shape(reprs["cls_repr"]) == (bsz, embed_dim)
+    assert ops.shape(reprs["encoder_rep"]) == (bsz, seq_len, embed_dim)
+
+
+def test_unimol_conf_gen_and_docking():
+    bsz = 2
+    seq_len = 6
+    vocab_size = 64
+    embed_dim = 64
+    heads = 4
+
+    conf_model = UniMolConfGenModel(
+        vocab_size=vocab_size,
+        encoder_layers=2,
+        encoder_embed_dim=embed_dim,
+        encoder_ffn_embed_dim=128,
+        encoder_attention_heads=heads,
+        num_kernel=32,
+    )
+    tokens = ops.convert_to_tensor(np.random.randint(1, vocab_size, (bsz, seq_len)), dtype="int32")
+    coord = ops.convert_to_tensor(np.random.randn(bsz, seq_len, 3).astype("float32"))
+
+    updated_coord, pred_dist = conf_model(tokens, src_coord=coord)
+    assert ops.shape(updated_coord) == (bsz, seq_len, 3)
+    assert ops.shape(pred_dist) == (bsz, seq_len, seq_len)
+
+    docking_model = UniMolDockingModel(
+        vocab_size=vocab_size,
+        encoder_layers=2,
+        encoder_embed_dim=embed_dim,
+        encoder_ffn_embed_dim=128,
+        encoder_attention_heads=heads,
+        num_kernel=32,
+    )
+    pose, pred_dist2 = docking_model(tokens, src_coord=coord)
+    assert ops.shape(pose) == (bsz, seq_len, 3)
+    assert ops.shape(pred_dist2) == (bsz, seq_len, seq_len)
+
+
+def test_unimol_weight_loading():
+    vocab_size = 32
+    embed_dim = 32
+    heads = 4
+    layers_num = 2
+
+    model = UniMolModel(
+        output_dim=2,
+        vocab_size=vocab_size,
+        encoder_layers=layers_num,
+        encoder_embed_dim=embed_dim,
+        encoder_ffn_embed_dim=64,
+        encoder_attention_heads=heads,
+        num_kernel=16,
+    )
+    model.build(None)
+
+    # Create synthetic PyTorch state_dict
+    state_dict = {
+        "embed_tokens.weight": torch.randn(vocab_size, embed_dim),
+        "gbf.means.weight": torch.randn(1, 16),
+        "gbf.stds.weight": torch.randn(1, 16),
+        "gbf.mul.weight": torch.randn(1024, 1),
+        "gbf.bias.weight": torch.randn(1024, 1),
+        "gbf_proj.linear1.weight": torch.randn(16, 16),
+        "gbf_proj.linear1.bias": torch.randn(16),
+        "gbf_proj.linear2.weight": torch.randn(heads, 16),
+        "gbf_proj.linear2.bias": torch.randn(heads),
+        "encoder.emb_layer_norm.weight": torch.ones(embed_dim),
+        "encoder.emb_layer_norm.bias": torch.zeros(embed_dim),
+        "classification_head.out_proj.weight": torch.randn(2, embed_dim),
+        "classification_head.out_proj.bias": torch.randn(2),
+    }
+
+    for i in range(layers_num):
+        prefix = f"encoder.layers.{i}"
+        state_dict[f"{prefix}.self_attn.in_proj.weight"] = torch.randn(embed_dim * 3, embed_dim)
+        state_dict[f"{prefix}.self_attn.in_proj.bias"] = torch.randn(embed_dim * 3)
+        state_dict[f"{prefix}.self_attn.out_proj.weight"] = torch.randn(embed_dim, embed_dim)
+        state_dict[f"{prefix}.self_attn.out_proj.bias"] = torch.randn(embed_dim)
+        state_dict[f"{prefix}.self_attn_layer_norm.weight"] = torch.ones(embed_dim)
+        state_dict[f"{prefix}.self_attn_layer_norm.bias"] = torch.zeros(embed_dim)
+        state_dict[f"{prefix}.fc1.weight"] = torch.randn(64, embed_dim)
+        state_dict[f"{prefix}.fc1.bias"] = torch.randn(64)
+        state_dict[f"{prefix}.fc2.weight"] = torch.randn(embed_dim, 64)
+        state_dict[f"{prefix}.fc2.bias"] = torch.randn(embed_dim)
+        state_dict[f"{prefix}.final_layer_norm.weight"] = torch.ones(embed_dim)
+        state_dict[f"{prefix}.final_layer_norm.bias"] = torch.zeros(embed_dim)
+
+    with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp:
+        tmp_path = tmp.name
+        torch.save({"model": state_dict}, tmp_path)
+
+    try:
+        load_unimol_weights(model, checkpoint_path=tmp_path, download=False)
+        # Check assigned weight matches
+        embed_w = ops.convert_to_numpy(model.embed_tokens.embeddings)
+        np.testing.assert_allclose(embed_w, state_dict["embed_tokens.weight"].numpy(), atol=1e-5)
+    finally:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+

k3_node/models/test_unimol2.py

@@ -0,0 +1,114 @@
+import os
+import tempfile
+import numpy as np
+import pytest
+import torch
+import keras
+from keras import ops
+
+from k3_node.models import (
+    UniMol2Model,
+    UniMol2AtomFeature,
+    UniMol2EdgeFeature,
+    UniMol2SE3Kernel,
+    UniMol2MovementHead,
+    download_unimol2_checkpoint,
+    load_unimol2_weights,
+)
+
+
+def test_unimol2_features():
+    bsz, seq_len = 2, 5
+    embed_dim = 32
+    pair_dim = 16
+
+    atom_feature = UniMol2AtomFeature(num_atom=128, embed_dim=embed_dim)
+    tokens = ops.convert_to_tensor(np.random.randint(0, 128, (bsz, seq_len)), dtype="int32")
+    x = atom_feature(tokens)
+    assert ops.shape(x) == (bsz, seq_len, embed_dim)
+
+    edge_feature = UniMol2EdgeFeature(num_edge=32, pair_dim=pair_dim)
+    edge_types = ops.convert_to_tensor(np.random.randint(0, 32, (bsz, seq_len, seq_len)), dtype="int32")
+    pair = edge_feature(edge_types)
+    assert ops.shape(pair) == (bsz, seq_len, seq_len, pair_dim)
+
+    se3 = UniMol2SE3Kernel(num_kernel=32, pair_dim=pair_dim)
+    dist = ops.convert_to_tensor(np.random.uniform(0.5, 5.0, (bsz, seq_len, seq_len)).astype("float32"))
+    pair_geo = se3(dist)
+    assert ops.shape(pair_geo) == (bsz, seq_len, seq_len, pair_dim)
+
+
+def test_unimol2_movement_head():
+    bsz, seq_len = 2, 6
+    pair_dim = 16
+    coords = ops.convert_to_tensor(np.random.randn(bsz, seq_len, 3).astype("float32"))
+    pair = ops.convert_to_tensor(np.random.randn(bsz, seq_len, seq_len, pair_dim).astype("float32"))
+
+    head = UniMol2MovementHead(pair_dim=pair_dim, hidden_dim=32)
+    updated = head(coords, pair)
+    assert ops.shape(updated) == (bsz, seq_len, 3)
+
+
+def test_unimol2_model_forward():
+    bsz, seq_len = 2, 6
+    embed_dim = 32
+    pair_dim = 16
+    num_heads = 4
+    num_layers = 2
+
+    model = UniMol2Model(
+        model_size="84m",
+        output_dim=2,
+        num_encoder_layers=num_layers,
+        encoder_embed_dim=embed_dim,
+        num_attention_heads=num_heads,
+        pair_embed_dim=pair_dim,
+        ffn_embedding_dim=64,
+        pair_hidden_dim=8,
+    )
+
+    atom_types = ops.convert_to_tensor(np.random.randint(0, 128, (bsz, seq_len)), dtype="int32")
+    coords = ops.convert_to_tensor(np.random.randn(bsz, seq_len, 3).astype("float32"))
+
+    # Logits
+    logits = model(atom_types, coords=coords)
+    assert ops.shape(logits) == (bsz, 2)
+
+    # Return coords and representations
+    logits, new_coords, x, pair = model(atom_types, coords=coords, return_coords=True, return_repr=True)
+    assert ops.shape(logits) == (bsz, 2)
+    assert ops.shape(new_coords) == (bsz, seq_len, 3)
+    assert ops.shape(x) == (bsz, seq_len, embed_dim)
+    assert ops.shape(pair) == (bsz, seq_len, seq_len, pair_dim)
+
+
+def test_unimol2_weight_loading():
+    model = UniMol2Model(
+        model_size="84m",
+        output_dim=2,
+        num_encoder_layers=1,
+        encoder_embed_dim=16,
+        num_attention_heads=2,
+        pair_embed_dim=8,
+        ffn_embedding_dim=16,
+        pair_hidden_dim=4,
+    )
+    model.build(None)
+
+    state_dict = {
+        "atom_feature.atom_embed.weight": torch.randn(512, 16),
+        "edge_feature.edge_embed.weight": torch.randn(64, 8),
+    }
+
+    with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp:
+        tmp_path = tmp.name
+        torch.save({"model": state_dict}, tmp_path)
+
+    try:
+        load_unimol2_weights(model, checkpoint_path=tmp_path, download=False)
+        w = ops.convert_to_numpy(model.atom_feature.atom_embed.embeddings)
+        np.testing.assert_allclose(w, state_dict["atom_feature.atom_embed.weight"].numpy(), atol=1e-5)
+    finally:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+