k3-node 1.0.0__py3-none-any.whl

k3_node/transforms/test_random_link_split.py

@@ -0,0 +1,45 @@
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data
+from k3_node.transforms import RandomLinkSplit
+
+
+def _graph():
+    rng = np.random.default_rng(0)
+    edge_index = rng.integers(0, 60, size=(2, 200))
+    return Data(x=rng.random((60, 4)).astype("float32"), edge_index=edge_index, num_nodes=60)
+
+
+def _np(x):
+    return np.asarray(ops.convert_to_numpy(x))
+
+
+def test_edge_labels_include_negatives():
+    train, val, test = RandomLinkSplit(num_val=0.1, num_test=0.2)(_graph())
+    for split, n_pos in ((train, 140), (val, 20), (test, 40)):
+        label = _np(split.edge_label)
+        assert label.dtype == np.float32
+        assert label.sum() == n_pos and (label == 0).sum() == n_pos  # one negative per positive
+        assert _np(split.edge_label_index).shape == (2, 2 * n_pos)
+    # message passing edges: training edges for train/val, plus validation edges for test
+    assert _np(train.edge_index).shape[1] == 140 and _np(test.edge_index).shape[1] == 160
+
+
+def test_no_train_negatives_and_ratio():
+    train, val, _ = RandomLinkSplit(num_val=0.1, num_test=0.2, add_negative_train_samples=False,
+                                    neg_sampling_ratio=2.0)(_graph())
+    assert (_np(train.edge_label) == 0).sum() == 0
+    assert (_np(val.edge_label) == 0).sum() == 40
+
+
+def test_split_labels_and_undirected():
+    edge_index = _np(_graph().edge_index)
+    edge_index = np.unique(np.sort(edge_index, axis=0), axis=1)
+    edge_index = edge_index[:, edge_index[0] != edge_index[1]]
+    data = Data(edge_index=np.concatenate([edge_index, edge_index[::-1]], axis=1), num_nodes=60)
+    train, val, test = RandomLinkSplit(num_val=0.1, num_test=0.2, is_undirected=True, split_labels=True)(data)
+    n = edge_index.shape[1]
+    assert _np(val.pos_edge_label_index).shape[1] == int(0.1 * n)
+    assert _np(val.neg_edge_label_index).shape[1] == int(0.1 * n)
+    assert _np(train.edge_index).shape[1] == 2 * _np(train.pos_edge_label_index).shape[1]

k3_node/transforms/test_spatial_transforms.py

@@ -0,0 +1,65 @@
+"""Spatial transforms on every backend: inputs may be NumPy arrays or backend tensors."""
+import numpy as np
+import pytest
+from keras import ops
+
+import k3_node.transforms as T
+from k3_node.data import Data
+
+
+def _data(as_tensor):
+    conv = ops.convert_to_tensor if as_tensor else np.asarray
+    pos = conv(np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype="float32"))
+    edge_index = conv(np.array([[0, 1, 1, 2], [1, 0, 2, 1]], dtype="int64"))
+    face = conv(np.array([[0], [1], [2]], dtype="int64"))
+    return Data(pos=pos, edge_index=edge_index, face=face, num_nodes=3)
+
+
+@pytest.mark.parametrize("as_tensor", [False, True])
+def test_spatial_transforms_all_backends(as_tensor):
+    data = _data(as_tensor)
+    np_ = ops.convert_to_numpy
+
+    assert tuple(T.Distance(norm=False)(data).edge_attr.shape) == (4, 1)
+    cart = T.Cartesian(norm=False)(data).edge_attr
+    assert tuple(cart.shape) == (4, 3)
+    np.testing.assert_allclose(np_(cart)[0], [-1.0, 0.0, 0.0])
+    assert str(np_(cart).dtype) == "float32"
+    assert tuple(T.Cartesian()(T.Distance()(data)).edge_attr.shape) == (4, 4)  # cat=True appends
+    assert tuple(T.LocalCartesian()(data).edge_attr.shape) == (4, 3)
+
+    data_2d = Data(pos=data.pos[:, :2], edge_index=data.edge_index)
+    assert tuple(T.Polar()(data_2d).edge_attr.shape) == (4, 2)
+    assert tuple(T.Spherical()(data).edge_attr.shape) == (4, 3)
+
+    np.testing.assert_allclose(np_(T.Center()(data).pos).mean(0), np.zeros(3), atol=1e-6)
+    assert np.abs(np_(T.NormalizeScale()(data).pos)).max() <= 1.0
+    assert tuple(T.NormalizeRotation(sort=True)(data).pos.shape) == (3, 3)
+
+    assert tuple(T.RandomJitter(0.1)(data).pos.shape) == (3, 3)
+    np.testing.assert_allclose(np_(T.RandomFlip(axis=0, p=1.0)(data).pos)[:, 0], -np_(data.pos)[:, 0])
+    np.testing.assert_allclose(np_(T.RandomScale((1.5, 1.5))(data).pos), np_(data.pos) * 1.5)
+    rot = T.RandomRotate(degrees=(90, 90), axis=2)(data).pos
+    np.testing.assert_allclose(np.linalg.norm(np_(rot), axis=1), np.linalg.norm(np_(data.pos), axis=1), atol=1e-6)
+    assert tuple(T.RandomShear(0.1)(data).pos.shape) == (3, 3)
+
+    f2e = T.FaceToEdge()(data)
+    assert f2e.face is None and tuple(f2e.edge_index.shape) == (2, 6)
+
+    mesh = T.GenerateMeshNormals()(data)
+    np.testing.assert_allclose(np.abs(np_(mesh.norm)[:, 2]), np.ones(3), atol=1e-6)
+    ppf = T.PointPairFeatures()(Data(pos=data.pos, norm=mesh.norm, edge_index=data.edge_index))
+    assert tuple(ppf.edge_attr.shape) == (4, 4)
+
+    sampled = T.SamplePoints(16, include_normals=True)(data)
+    assert tuple(sampled.pos.shape) == (16, 3) and tuple(sampled.normal.shape) == (16, 3)
+    assert np.all(np.abs(np_(sampled.pos)[:, 2]) < 1e-6)  # all samples lie on the triangle's plane
+
+    fixed = T.FixedPoints(2)(data)
+    assert tuple(fixed.pos.shape) == (2, 3) and fixed.num_nodes == 2
+
+    grid = T.GridSampling(size=0.5)(Data(pos=data.pos))
+    assert grid.num_nodes == 3
+
+    square = Data(pos=np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], dtype="float32"))
+    assert tuple(T.Delaunay()(square).face.shape) == (3, 2)

k3_node/transforms/test_transforms.py

@@ -0,0 +1,253 @@
+import math
+import numpy as np
+import pytest
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+if torch is None:
+    pytest.skip("PyTorch is required for transforms unit tests", allow_module_level=True)
+
+import k3_node.transforms as T
+from k3_node.data import Data, HeteroData
+
+
+def test_compose():
+    data = Data(x=torch.ones((4, 2)), edge_index=torch.tensor([[0, 1], [1, 2]]))
+    transform = T.Compose([
+        T.Constant(value=2.0, cat=False),
+        T.AddSelfLoops(),
+    ])
+    out = transform(data)
+    assert out.x.shape == (4, 1)
+    assert out.x[0, 0].item() == 2.0
+    assert out.edge_index.shape[1] == 2 + 4
+
+
+def test_compose_filters():
+    data1 = Data(num_nodes=5)
+    data2 = Data(num_nodes=2)
+    filt = T.ComposeFilters([lambda d: d.num_nodes > 3])
+    assert filt(data1) is True
+    assert filt(data2) is False
+
+
+def test_general_transforms():
+    x = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
+    edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long)
+    data = Data(x=x, edge_index=edge_index)
+
+    # Constant
+    d_const = T.Constant(value=5.0, cat=True)(data)
+    assert d_const.x.shape == (2, 3)
+
+    # NormalizeFeatures
+    d_norm = T.NormalizeFeatures()(data)
+    assert torch.allclose(d_norm.x.sum(dim=-1), torch.tensor([1.0, 1.0]))
+
+    # SVDFeatureReduction
+    d_svd = T.SVDFeatureReduction(out_channels=1)(data)
+    assert d_svd.x.shape == (2, 1)
+
+    # IndexToMask & MaskToIndex
+    data_idx = Data(num_nodes=4, train_index=torch.tensor([0, 2]))
+    d_mask = T.IndexToMask(attrs='train_index')(data_idx)
+    assert hasattr(d_mask, 'train_mask')
+    assert d_mask.train_mask.tolist() == [True, False, True, False]
+    d_idx2 = T.MaskToIndex(attrs='train_mask')(d_mask)
+    assert hasattr(d_idx2, 'train_index')
+    assert d_idx2.train_index.tolist() == [0, 2]
+
+    # Pad
+    d_pad = T.Pad(max_num_nodes=5)(data)
+    assert d_pad.x.shape == (5, 2)
+    assert d_pad.num_nodes == 5
+
+    # RandomNodeSplit
+    data_nodes = Data(num_nodes=10, y=torch.tensor([0, 1] * 5))
+    d_split = T.RandomNodeSplit(split='train_rest', num_val=2, num_test=2)(data_nodes)
+    assert d_split.train_mask.sum() == 6
+    assert d_split.val_mask.sum() == 2
+    assert d_split.test_mask.sum() == 2
+
+    # RemoveTrainingClasses
+    data_cls = Data(y=torch.tensor([0, 1, 2, 0, 1, 2]))
+    d_rem = T.RemoveTrainingClasses(classes=[0])(data_cls)
+    assert (d_rem.y == 0).sum() == 0
+
+    # ToSparseTensor
+    d_adj = T.ToSparseTensor()(data)
+    assert hasattr(d_adj, 'adj_t')
+
+
+def test_graph_transforms():
+    edge_index = torch.tensor([[0, 1, 1], [1, 0, 2]], dtype=torch.long)
+    x = torch.randn(3, 4)
+    data = Data(x=x, edge_index=edge_index, num_nodes=3)
+
+    # ToUndirected
+    d_undir = T.ToUndirected()(data)
+    assert d_undir.is_undirected()
+
+    # OneHotDegree
+    d_ohd = T.OneHotDegree(max_degree=3)(data)
+    assert d_ohd.x.shape[0] == 3
+
+    # TargetIndegree
+    d_tid = T.TargetIndegree()(data)
+    assert d_tid.edge_attr is not None
+
+    # LocalDegreeProfile
+    d_ldp = T.LocalDegreeProfile()(data)
+    assert d_ldp.x.shape == (3, 4 + 5)
+
+    # AddSelfLoops & RemoveSelfLoops
+    d_sl = T.AddSelfLoops()(data)
+    assert d_sl.edge_index.shape[1] == 3 + 3
+    d_rem_sl = T.RemoveSelfLoops()(d_sl)
+    assert d_rem_sl.edge_index.shape[1] == 3
+
+    # RemoveIsolatedNodes
+    data_iso = Data(edge_index=torch.tensor([[0], [1]]), num_nodes=3, x=torch.randn(3, 2))
+    d_no_iso = T.RemoveIsolatedNodes()(data_iso)
+    assert d_no_iso.num_nodes == 2
+
+    # RemoveDuplicatedEdges
+    data_dup = Data(edge_index=torch.tensor([[0, 0], [1, 1]]), num_nodes=2)
+    d_no_dup = T.RemoveDuplicatedEdges()(data_dup)
+    assert d_no_dup.edge_index.shape[1] == 1
+
+    # KNNGraph & RadiusGraph
+    data_pos = Data(pos=torch.tensor([[0.0, 0.0], [0.1, 0.0], [2.0, 2.0]]))
+    d_knn = T.KNNGraph(k=1)(data_pos)
+    assert d_knn.edge_index is not None
+    d_rad = T.RadiusGraph(r=0.5)(data_pos)
+    assert d_rad.edge_index is not None
+
+    # TwoHop
+    d_2hop = T.TwoHop()(data)
+    assert d_2hop.edge_index.shape[1] >= 3
+
+    # LineGraph
+    d_line = T.LineGraph()(data)
+    assert d_line.num_nodes == 3
+
+    # LaplacianLambdaMax
+    d_lap = T.LaplacianLambdaMax()(data)
+    assert hasattr(d_lap, 'lambda_max')
+
+    # GDC
+    d_gdc = T.GDC(self_loop_weight=1.0, normalization_in='sym', normalization_out='sym',
+                  diffusion_kwargs=dict(method='ppr', alpha=0.15))(data)
+    assert d_gdc.edge_index is not None
+
+    # SIGN
+    d_sign = T.SIGN(K=2)(data)
+    assert hasattr(d_sign, 'x1') and hasattr(d_sign, 'x2')
+
+    # GCNNorm
+    d_gcn = T.GCNNorm()(data)
+    assert d_gcn.edge_weight is not None
+
+    # VirtualNode
+    d_vn = T.VirtualNode()(data)
+    assert d_vn.num_nodes == 4
+
+    # AddLaplacianEigenvectorPE
+    d_lpe = T.AddLaplacianEigenvectorPE(k=2)(data)
+    assert hasattr(d_lpe, 'laplacian_eigenvector_pe')
+
+    # AddRandomWalkPE
+    d_rwpe = T.AddRandomWalkPE(walk_length=4)(data)
+    assert hasattr(d_rwpe, 'random_walk_pe')
+
+    # FeaturePropagation
+    d_fp = T.FeaturePropagation(missing_mask=torch.tensor([True, False, False]), num_iterations=5)(data)
+    assert d_fp.x.shape == (3, 4)
+
+    # HalfHop
+    d_hh = T.HalfHop(alpha=0.5)(data)
+    assert d_hh.num_nodes > 3
+
+
+def test_spatial_transforms():
+    pos = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=torch.float32)
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
+    face = torch.tensor([[0], [1], [2]], dtype=torch.long)
+    data = Data(pos=pos, edge_index=edge_index, face=face, num_nodes=3)
+
+    # Distance
+    d_dist = T.Distance(norm=False)(data)
+    assert d_dist.edge_attr.shape == (4, 1)
+
+    # Cartesian & LocalCartesian
+    d_cart = T.Cartesian(norm=False)(data)
+    assert d_cart.edge_attr.shape == (4, 3)
+    d_lcart = T.LocalCartesian(norm=False)(data)
+    assert d_lcart.edge_attr.shape == (4, 3)
+
+    # Polar & Spherical
+    data_2d = Data(pos=pos[:, :2], edge_index=edge_index)
+    d_pol = T.Polar(norm=False)(data_2d)
+    assert d_pol.edge_attr.shape == (4, 2)
+    d_sph = T.Spherical(norm=False)(data)
+    assert d_sph.edge_attr.shape == (4, 3)
+
+    # Center, NormalizeScale, NormalizeRotation
+    d_cen = T.Center()(data)
+    assert torch.allclose(d_cen.pos.mean(dim=0), torch.zeros(3), atol=1e-5)
+
+    d_nscale = T.NormalizeScale()(data)
+    assert d_nscale.pos.abs().max() <= 1.0
+
+    d_nrot = T.NormalizeRotation()(data)
+    assert d_nrot.pos.shape == (3, 3)
+
+    # RandomJitter, RandomFlip, RandomScale, RandomRotate, RandomShear
+    d_jit = T.RandomJitter(translate=0.1)(data)
+    assert d_jit.pos.shape == (3, 3)
+
+    d_flip = T.RandomFlip(axis=0, p=1.0)(data)
+    assert torch.allclose(d_flip.pos[:, 0], -data.pos[:, 0])
+
+    d_scale = T.RandomScale(scales=(1.5, 1.5))(data)
+    assert torch.allclose(d_scale.pos, data.pos * 1.5)
+
+    d_rot = T.RandomRotate(degrees=90, axis=0)(data)
+    assert d_rot.pos.shape == (3, 3)
+
+    d_shear = T.RandomShear(shear=0.1)(data)
+    assert d_shear.pos.shape == (3, 3)
+
+    # FaceToEdge
+    d_f2e = T.FaceToEdge(remove_faces=True)(data)
+    assert d_f2e.face is None
+    assert d_f2e.edge_index is not None
+
+    # GenerateMeshNormals
+    d_mesh = T.GenerateMeshNormals()(data)
+    assert hasattr(d_mesh, 'norm')
+    assert d_mesh.norm.shape == (3, 3)
+
+    # PointPairFeatures
+    data_ppf = Data(pos=pos, norm=d_mesh.norm, edge_index=edge_index)
+    d_ppf = T.PointPairFeatures()(data_ppf)
+    assert d_ppf.edge_attr.shape == (4, 4)
+
+    # Delaunay
+    data_pts = Data(pos=torch.tensor([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]))
+    d_del = T.Delaunay()(data_pts)
+    assert hasattr(d_del, 'face')
+
+    # FixedPoints
+    d_fixed = T.FixedPoints(num=2)(data)
+    assert d_fixed.pos.shape[0] == 2
+    assert d_fixed.num_nodes == 2
+
+    # GridSampling
+    data_cloud = Data(pos=pos)
+    d_grid = T.GridSampling(size=0.5)(data_cloud)
+    assert d_grid.pos is not None
+    assert d_grid.num_nodes > 0

k3_node/transforms/utils.py

@@ -0,0 +1,102 @@
+import numpy as np
+from keras import ops
+from typing import Any, List, Optional, Tuple, Union
+
+try:
+    import torch
+    _HAS_TORCH = True
+except ImportError:
+    _HAS_TORCH = False
+
+
+def is_torch_tensor(x: Any) -> bool:
+    return _HAS_TORCH and isinstance(x, torch.Tensor)
+
+
+def to_numpy(x: Any) -> np.ndarray:
+    if x is None:
+        return None
+    if isinstance(x, np.ndarray):
+        return x
+    from k3_node.ops.host import to_numpy as host_to_numpy
+
+    if is_torch_tensor(x) and x.device.type != "meta":
+        return x.detach().cpu().numpy()
+    return host_to_numpy(x)
+
+
+def as_tensor(x: Any, dtype: Optional[Any] = None, like: Optional[Any] = None) -> Any:
+    if x is None:
+        return None
+    if like is not None and is_torch_tensor(like):
+        if isinstance(dtype, str):
+            dtype_map = {
+                "float32": torch.float32,
+                "float64": torch.float64,
+                "float16": torch.float16,
+                "int64": torch.int64,
+                "int32": torch.int32,
+                "int16": torch.int16,
+                "int8": torch.int8,
+                "uint8": torch.uint8,
+                "bool": torch.bool,
+            }
+            target_dtype = dtype_map.get(dtype, getattr(torch, dtype, None))
+        else:
+            target_dtype = dtype or like.dtype
+        target_device = like.device
+        if is_torch_tensor(x):
+            return x.to(device=target_device, dtype=target_dtype)
+        np_arr = to_numpy(x)
+        return torch.as_tensor(np_arr, dtype=target_dtype, device=target_device)
+    if is_torch_tensor(x):
+        if dtype is not None:
+            if isinstance(dtype, str):
+                target_dtype = getattr(torch, dtype, None) or getattr(torch, dtype.replace("torch.", ""), None)
+            else:
+                target_dtype = dtype
+            return x.to(dtype=target_dtype)
+        return x
+    if isinstance(x, np.ndarray) and like is not None and isinstance(like, np.ndarray):
+        return np.asarray(x, dtype=dtype or like.dtype)
+    return ops.convert_to_tensor(x, dtype=dtype)
+
+
+def match_tensor(new_val: Any, reference: Any, dtype: Optional[Any] = None) -> Any:
+    if reference is None:
+        if is_torch_tensor(new_val):
+            return new_val
+        return ops.convert_to_tensor(new_val, dtype=dtype)
+    return as_tensor(new_val, dtype=dtype, like=reference)
+
+
+def to_undirected(
+    edge_index: Any,
+    edge_attr: Optional[Any] = None,
+    num_nodes: Optional[int] = None,
+    reduce: str = "add",
+) -> Union[Any, Tuple[Any, Any]]:
+    from k3_node.utils.graph import coalesce
+
+    ei_np = to_numpy(edge_index)
+    if ei_np.shape[1] == 0:
+        if edge_attr is None:
+            return edge_index
+        return edge_index, edge_attr
+
+    row, col = ei_np[0], ei_np[1]
+    rev_row = np.concatenate([row, col], axis=0)
+    rev_col = np.concatenate([col, row], axis=0)
+    rev_edge_index = np.stack([rev_row, rev_col], axis=0)
+
+    if edge_attr is not None:
+        attr_np = to_numpy(edge_attr)
+        if attr_np.ndim == 1:
+            rev_attr = np.concatenate([attr_np, attr_np], axis=0)
+        else:
+            rev_attr = np.concatenate([attr_np, attr_np], axis=0)
+        out_ei, out_attr = coalesce(rev_edge_index, rev_attr, num_nodes=num_nodes, reduce=reduce)
+        return match_tensor(out_ei, edge_index), match_tensor(out_attr, edge_attr)
+    else:
+        out_ei, _ = coalesce(rev_edge_index, None, num_nodes=num_nodes, reduce=reduce)
+        return match_tensor(out_ei, edge_index)

k3_node/utils/__init__.py

@@ -0,0 +1,5 @@
+from .keras import *
+from .backend_import import *  # noqa
+from .graph import *
+from .random import *
+from .smiles import *

k3_node/utils/backend_import.py

@@ -0,0 +1,12 @@
+from keras import backend
+
+if backend.backend() == "tensorflow":
+    import tensorflow as tf
+elif backend.backend() == "torch":
+    import torch
+elif backend.backend() == "jax":
+    import jax.numpy as jnp
+elif backend.backend() == "numpy":
+    import numpy as np
+else:
+    raise NotImplementedError