1 files changed, 100 insertions, 0 deletions
diff --git a/kaft_mvp/data.py b/kaft_mvp/data.py
new file mode 100644
index 0000000..2f25d80
--- /dev/null
+++ b/kaft_mvp/data.py
@@ -0,0 +1,100 @@
+"""Dataset and graph-operator utilities for the MVP."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+from urllib.request import urlopen
+
+import torch
+import torch_geometric.transforms as transforms
+from torch_geometric.datasets import Planetoid
+
+
+PLANETOID_FILES = (
+    "x",
+    "tx",
+    "allx",
+    "y",
+    "ty",
+    "ally",
+    "graph",
+    "test.index",
+)
+
+
+def _ensure_planetoid_raw(root: str | Path, name: str) -> None:
+    """Pre-fetch from raw.githubusercontent.com to avoid a flaky redirect."""
+    raw_dir = Path(root) / name / "raw"
+    raw_dir.mkdir(parents=True, exist_ok=True)
+    lower_name = name.lower()
+    base_url = (
+        "https://raw.githubusercontent.com/kimiyoung/"
+        "planetoid/master/data"
+    )
+    for suffix in PLANETOID_FILES:
+        filename = f"ind.{lower_name}.{suffix}"
+        destination = raw_dir / filename
+        if destination.exists() and destination.stat().st_size > 0:
+            continue
+        temporary = destination.with_suffix(destination.suffix + ".part")
+        with urlopen(f"{base_url}/{filename}", timeout=60) as response:
+            temporary.write_bytes(response.read())
+        temporary.replace(destination)
+
+
+def spmm(matrix: torch.Tensor, dense: torch.Tensor) -> torch.Tensor:
+    """Multiply a native PyTorch sparse matrix by a dense matrix."""
+    return torch.sparse.mm(matrix, dense)
+
+
+def normalized_adjacency(
+    edge_index: torch.Tensor,
+    num_nodes: int,
+) -> torch.Tensor:
+    """Return A-hat = D^-1/2 (A + I) D^-1/2 as a coalesced COO tensor."""
+    row, col = edge_index
+    loops = torch.arange(num_nodes, device=edge_index.device)
+    row = torch.cat((row, loops))
+    col = torch.cat((col, loops))
+
+    degree = torch.zeros(num_nodes, device=edge_index.device)
+    degree.scatter_add_(0, row, torch.ones_like(row, dtype=torch.float32))
+    inverse_sqrt = degree.pow(-0.5)
+    inverse_sqrt.masked_fill_(torch.isinf(inverse_sqrt), 0.0)
+    values = inverse_sqrt[row] * inverse_sqrt[col]
+
+    return torch.sparse_coo_tensor(
+        torch.stack((row, col)),
+        values,
+        (num_nodes, num_nodes),
+    ).coalesce()
+
+
+def load_cora(
+    root: str | Path = "data",
+    device: str | torch.device = "cpu",
+) -> dict[str, Any]:
+    """Download Cora and return the tensors used by both training rules."""
+    _ensure_planetoid_raw(root, "Cora")
+    dataset = Planetoid(
+        root=str(root),
+        name="Cora",
+        transform=transforms.NormalizeFeatures(),
+    )
+    graph = dataset[0]
+    adjacency = normalized_adjacency(
+        graph.edge_index,
+        graph.num_nodes,
+    ).to(device)
+    return {
+        "x": graph.x.to(device),
+        "y": graph.y.to(device),
+        "adjacency": adjacency,
+        "train_mask": graph.train_mask.to(device),
+        "val_mask": graph.val_mask.to(device),
+        "test_mask": graph.test_mask.to(device),
+        "num_nodes": graph.num_nodes,
+        "num_features": dataset.num_features,
+        "num_classes": dataset.num_classes,
+    }