【PyG】PyTorch Geometricのインストール方法から利用方法まで解説

import torch
print(torch.__version__)        # 例: 2.1.0
print(torch.version.cuda)       # 例: 11.8 (CPU環境では None)

pip install torch_geometric

# torch-scatter, torch-sparse をインストールするには PyTorch/CUDA バージョンを指定
# 例: PyTorch 2.1.0, CUDA 11.8 の場合
pip install pyg_lib torch_scatter torch_sparse -f https://data.pyg.org/whl/torch-2.1.0+cu118.html

conda install pyg -c pyg

import torch_geometric
print(torch_geometric.__version__)  # 例: 2.5.3

import torch

# 送信元ノードのリスト（edge_index[0]）
src = [0, 0, 1, 2, 1]
# 受信先ノードのリスト（edge_index[1]）
dst = [1, 2, 3, 3, 2]

edge_index = torch.tensor([src, dst], dtype=torch.long)
print(edge_index.shape)  # torch.Size([2, 5])

from torch_geometric.data import Data

# ノード特徴量: 3ノード、各2次元
x = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
# 辺: 0→1, 0→2, 2→1
edge_index = torch.tensor([[0, 0, 2], [1, 2, 1]], dtype=torch.long)
# ノードラベル
y = torch.tensor([0, 1, 2])

data = Data(x=x, edge_index=edge_index, y=y)
print(data)
# Data(x=[3, 2], edge_index=[2, 3], y=[3])

# 便利なプロパティ
print(data.num_nodes)  # 3
print(data.num_edges)  # 3
print(data.is_directed())  # True

import networkx as nx
import matplotlib.pyplot as plt
from torch_geometric.utils import to_networkx

nxg = to_networkx(data, to_undirected=False)
nx.draw(nxg, with_labels=True, node_color="steelblue", font_color="white",
        node_size=800, arrows=True)
plt.title("PyG Data → NetworkX グラフ")
plt.show()

# 自己ループがないか、孤立ノードがないかを確認
print(data.has_self_loops())  # False
print(data.is_undirected())   # False（有向グラフ）

# edge_index の最大インデックスが num_nodes 未満か検証
assert data.edge_index.max() < data.num_nodes

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        # 2層のGCN
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        # 第1層: メッセージパッシング + ReLU + Dropout
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        # 第2層: メッセージパッシング（活性化なし）
        x = self.conv2(x, edge_index)
        return x  # ノード分類の場合、softmax は loss 計算時にかける

from torch_geometric.nn import GATConv

class GAT(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, heads=8):
        super().__init__()
        # マルチヘッドアテンション: heads 個を並列実行し結合
        self.conv1 = GATConv(in_channels, hidden_channels, heads=heads,
                             dropout=0.6)
        # 第2層は heads=1 で集約（または concat=False で平均）
        self.conv2 = GATConv(hidden_channels * heads, out_channels, heads=1,
                             concat=False, dropout=0.6)

    def forward(self, x, edge_index):
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        return x

from torch_geometric.nn import SAGEConv

class GraphSAGE(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv(in_channels, hidden_channels, aggr="mean")
        self.conv2 = SAGEConv(hidden_channels, out_channels, aggr="mean")

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

from torch_geometric.loader import DataLoader

# 複数のグラフからなるデータセットを仮定
graphs = [
    Data(x=torch.randn(3, 8), edge_index=torch.tensor([[0,1],[1,2]]), y=torch.tensor([0])),
    Data(x=torch.randn(5, 8), edge_index=torch.tensor([[0,1,2,3],[1,2,3,4]]), y=torch.tensor([1])),
    Data(x=torch.randn(4, 8), edge_index=torch.tensor([[0,1,2],[1,2,3]]), y=torch.tensor([0])),
]

loader = DataLoader(graphs, batch_size=2, shuffle=True)

for batch in loader:
    print(batch)
    # DataBatch(x=[8, 8], edge_index=[2, 3], y=[2], batch=[8], ptr=[3])
    # x=[8, 8] は バッチ内の全ノード数 x 特徴次元
    # batch=[8] は各ノードがグラフ0か1かを示すインデックス
    print(batch.batch)   # tensor([0, 0, 0, 1, 1, 1, 1, 1])
    print(batch.num_graphs)  # 2
    break

from torch_geometric.datasets import TUDataset

# MUTAG: 188グラフ、2クラス分類（化合物の変異原性）
dataset = TUDataset(root="/tmp/MUTAG", name="MUTAG")
print(f"グラフ数: {len(dataset)}")           # 188
print(f"クラス数: {dataset.num_classes}")      # 2
print(f"特徴次元: {dataset.num_node_features}") # 7

# 訓練・テストに分割
train_dataset = dataset[:150]
test_dataset  = dataset[150:]

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader  = DataLoader(test_dataset, batch_size=32, shuffle=False)

from torch_geometric.datasets import Planetoid

# データセットのダウンロード
dataset = Planetoid(root="/tmp/Cora", name="Cora")
data = dataset[0]  # グラフは1つ

print(f"ノード数: {data.num_nodes}")          # 2708
print(f"辺数: {data.num_edges}")             # 10556
print(f"特徴次元: {data.num_node_features}") # 1433
print(f"クラス数: {dataset.num_classes}")    # 7
print(f"訓練ノード数: {data.train_mask.sum()}") # 140
print(f"検証ノード数: {data.val_mask.sum()}")   # 500
print(f"テストノード数: {data.test_mask.sum()}") # 1000

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCNClassifier(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, dropout=0.5):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)
        self.dropout = dropout

    def forward(self, x, edge_index):
        # 第1層
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        # 第2層
        x = self.conv2(x, edge_index)
        return x  # 生のロジット

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = GCNClassifier(
    in_channels=dataset.num_node_features,  # 1433
    hidden_channels=64,
    out_channels=dataset.num_classes,       # 7
).to(device)

data = data.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    # 訓練ノードのみで損失計算
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

@torch.no_grad()
def evaluate():
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=-1)

    accs = {}
    for split, mask in [("train", data.train_mask),
                         ("val",   data.val_mask),
                         ("test",  data.test_mask)]:
        correct = pred[mask] == data.y[mask]
        accs[split] = correct.sum().item() / mask.sum().item()
    return accs

best_val_acc = 0
for epoch in range(1, 201):
    loss = train()
    accs = evaluate()
    if accs["val"] > best_val_acc:
        best_val_acc = accs["val"]
        best_test_acc = accs["test"]
    if epoch % 20 == 0:
        print(f"Epoch {epoch:03d}: Loss={loss:.4f} "
              f"Train={accs['train']:.4f} Val={accs['val']:.4f} Test={accs['test']:.4f}")

from torch_geometric.nn import GCNConv, JumpingKnowledge
import torch.nn as nn

class DeepGCN(torch.nn.Module):
    """Jumping Knowledge でスキップ接続を追加した深い GCN"""
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers=4):
        super().__init__()
        self.convs = nn.ModuleList()
        self.convs.append(GCNConv(in_channels, hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(GCNConv(hidden_channels, hidden_channels))
        self.convs.append(GCNConv(hidden_channels, hidden_channels))

        # JK: 全層の出力を max または concatenate で結合
        self.jk = JumpingKnowledge(mode="max")
        self.lin = nn.Linear(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        xs = []
        for conv in self.convs:
            x = F.relu(conv(x, edge_index))
            x = F.dropout(x, p=0.5, training=self.training)
            xs.append(x)
        # 各層の出力を結合して最終表現を得る
        out = self.jk(xs)
        return self.lin(out)

# PyG での GINConv の使い方
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GINConv, global_add_pool

class GIN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers=3):
        super().__init__()
        self.convs = nn.ModuleList()
        self.bns = nn.ModuleList()  # Batch Normalization

        for i in range(num_layers):
            in_ch = in_channels if i == 0 else hidden_channels
            # GINConv には任意の NN を渡せる（内部で MLP として使われる）
            mlp = nn.Sequential(
                nn.Linear(in_ch, hidden_channels),
                nn.BatchNorm1d(hidden_channels),
                nn.ReLU(),
                nn.Linear(hidden_channels, hidden_channels),
            )
            self.convs.append(GINConv(mlp, eps=0, train_eps=True))
            self.bns.append(nn.BatchNorm1d(hidden_channels))

        self.lin1 = nn.Linear(hidden_channels, hidden_channels)
        self.lin2 = nn.Linear(hidden_channels, out_channels)

    def forward(self, x, edge_index, batch):
        for conv, bn in zip(self.convs, self.bns):
            x = F.relu(bn(conv(x, edge_index)))

        # グラフ分類: 和集約（表現力を最大化）
        x = global_add_pool(x, batch)

        x = F.dropout(F.relu(self.lin1(x)), p=0.5, training=self.training)
        return self.lin2(x)

# 構文確認（PyG 環境で動作可能なコード）
model = GIN(in_channels=7, hidden_channels=64, out_channels=2)
print(model)
# GIN(
#   (convs): ModuleList(...)
#   (bns): ModuleList(...)
# )

from torch_geometric.loader import NeighborLoader
from torch_geometric.datasets import Planetoid

# 大規模グラフを想定した設定
dataset = Planetoid(root="/tmp/Cora", name="Cora")
data = dataset[0]

# NeighborLoader の設定
# num_neighbors=[25, 10]: 1層目で25近傍、2層目で10近傍をサンプリング
train_loader = NeighborLoader(
    data,
    num_neighbors=[25, 10],       # 各層でサンプリングする近傍数
    batch_size=1024,               # バッチごとのターゲットノード数
    input_nodes=data.train_mask,  # 訓練ノードだけをターゲットに
    shuffle=True,
)

# サブグラフのサイズを確認
for subgraph in train_loader:
    print(f"ターゲットノード数: {subgraph.batch_size}")         # 1024 (または以下)
    print(f"サブグラフのノード数: {subgraph.num_nodes}")        # サンプルした全ノード
    print(f"サブグラフの辺数: {subgraph.num_edges}")            # サブグラフ内の辺
    # subgraph.n_id には元グラフのノードIDが入っている
    print(f"ノードID範囲: {subgraph.n_id.min()} ~ {subgraph.n_id.max()}")
    break

import torch
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv

class GraphSAGELarge(torch.nn.Module):
    """大規模グラフ向け GraphSAGE"""
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv(in_channels, hidden_channels)
        self.conv2 = SAGEConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.5, training=self.training)
        return self.conv2(x, edge_index)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GraphSAGELarge(dataset.num_node_features, 256, dataset.num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=5e-4)

def train_large():
    model.train()
    total_loss = 0
    for subgraph in train_loader:
        subgraph = subgraph.to(device)
        optimizer.zero_grad()
        out = model(subgraph.x, subgraph.edge_index)
        # ターゲットノード (先頭 batch_size 個) だけで損失計算
        loss = F.cross_entropy(
            out[:subgraph.batch_size],
            subgraph.y[:subgraph.batch_size]
        )
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(train_loader)

# 注: この例は NeighborLoader の使用パターンを示す構文確認用コード。
# 実際の訓練は PyG + CUDA 環境で実行してください。

from torch_geometric.nn import GCNConv, global_mean_pool

class GraphClassifier(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.conv3 = GCNConv(hidden_channels, hidden_channels)
        self.lin1  = torch.nn.Linear(hidden_channels, hidden_channels)
        self.lin2  = torch.nn.Linear(hidden_channels, out_channels)

    def forward(self, x, edge_index, batch):
        # ノードレベルのメッセージパッシング
        x = F.relu(self.conv1(x, edge_index))
        x = F.relu(self.conv2(x, edge_index))
        x = F.relu(self.conv3(x, edge_index))

        # グラフレベルへの集約（batch テンソルがどのグラフかを示す）
        x = global_mean_pool(x, batch)  # [num_graphs, hidden_channels]

        # MLP 分類器
        x = F.relu(self.lin1(x))
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin2(x)
        return x

# MUTAG データセットでの例
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader

dataset = TUDataset(root="/tmp/MUTAG", name="MUTAG").shuffle()
train_loader = DataLoader(dataset[:150], batch_size=32, shuffle=True)
test_loader  = DataLoader(dataset[150:], batch_size=32, shuffle=False)

model = GraphClassifier(
    in_channels=dataset.num_node_features,  # 7
    hidden_channels=64,
    out_channels=dataset.num_classes,       # 2
)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

def train_graph():
    model.train()
    total_loss = 0
    for data in train_loader:
        optimizer.zero_grad()
        # forward に batch テンソルを渡す
        out = model(data.x, data.edge_index, data.batch)
        loss = F.cross_entropy(out, data.y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs
    return total_loss / len(train_loader.dataset)

@torch.no_grad()
def test_graph(loader):
    model.eval()
    correct = 0
    for data in loader:
        out = model(data.x, data.edge_index, data.batch)
        pred = out.argmax(dim=-1)
        correct += (pred == data.y).sum().item()
    return correct / len(loader.dataset)

# 注: このコードは PyG インストール環境での実行を想定。
# 実際の訓練結果は MUTAG で 70〜80% 前後の精度になることが多い。

from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class WeightedConv(MessagePassing):
    """辺の特徴量を重みとして使う加重平均集約"""

    def __init__(self, in_channels, out_channels):
        # aggr="add" で集約は和、update で正規化する
        super().__init__(aggr="add")
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index, edge_weight=None):
        # 自己ループを追加して自分自身の特徴も集約対象に
        edge_index, edge_weight = add_self_loops(
            edge_index, edge_weight, fill_value=1.0, num_nodes=x.size(0)
        )
        # 線形変換
        x = self.lin(x)
        # 次数の計算（正規化に使う）
        row, col = edge_index
        deg = degree(col, x.size(0), dtype=x.dtype)
        deg_inv = deg.pow(-1.0)
        deg_inv[deg_inv == float("inf")] = 0  # 孤立ノードの対処

        # propagate が message → aggregate → update を自動実行
        return self.propagate(edge_index, x=x,
                              edge_weight=edge_weight, deg_inv=deg_inv)

    def message(self, x_j, edge_weight, deg_inv_i):
        # 送信ノード j の特徴 x_j を辺重みで重み付け
        # deg_inv_i は受信ノード i の次数逆数（正規化用）
        return deg_inv_i.view(-1, 1) * edge_weight.view(-1, 1) * x_j

    def update(self, aggr_out):
        # 集約結果をそのまま返す（非線形変換を追加することも可）
        return aggr_out

# 使用例
conv = WeightedConv(in_channels=16, out_channels=32)
x = torch.randn(5, 16)  # 5ノード、16次元
edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]], dtype=torch.long)
edge_weight = torch.tensor([0.5, 0.8, 1.0, 0.3])

out = conv(x, edge_index, edge_weight)
print(out.shape)  # torch.Size([5, 32])

import torch_geometric.transforms as T
from torch_geometric.datasets import Planetoid

# Transforms をリストで組み合わせる（Compose で直列処理）
transform = T.Compose([
    T.NormalizeFeatures(),       # ノード特徴量を行ごとに L1 正規化
    T.ToUndirected(),            # 有向辺→無向辺（逆方向の辺を追加）
    T.AddSelfLoops(),            # 自己ループを追加（GCNConv では内部でも追加するが明示的に）
])

dataset = Planetoid(root="/tmp/Cora", name="Cora", transform=transform)
data = dataset[0]
print(data)
# Data(x=[2708, 1433], edge_index=[2, 13264], y=[2708], ...)
# 辺数が増えている（双方向化 + 自己ループ追加のため）

# pre_transform: 最初の1回だけ適用されディスクにキャッシュ
# transform: データアクセスのたびに適用（ランダム変換などに使う）
dataset = TUDataset(
    root="/tmp/MUTAG",
    name="MUTAG",
    pre_transform=T.NormalizeFeatures(),  # キャッシュされる
    transform=T.RandomFlip(axis=0),        # アクセスごとに適用
)

from torch_geometric.data import InMemoryDataset, Data, download_url
import torch
import os

class MyMoleculeDataset(InMemoryDataset):
    """独自の分子グラフデータセット（InMemoryDataset の実装例）"""

    def __init__(self, root, transform=None, pre_transform=None):
        # root: データのキャッシュ先ディレクトリ
        super().__init__(root, transform, pre_transform)
        # 前処理済みデータを読み込む
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        """生データのファイル名リスト（download で作られるファイル）"""
        return ["molecules.csv"]

    @property
    def processed_file_names(self):
        """処理済みデータのファイル名（存在すればダウンロード・処理をスキップ）"""
        return ["data.pt"]

    def download(self):
        """生データをダウンロードする（省略可）"""
        # download_url("https://example.com/molecules.csv", self.raw_dir)
        pass  # 既に raw_dir にデータがある場合はスキップ

    def process(self):
        """生データを読み込んで Data オブジェクトのリストを作成"""
        data_list = []

        # 例: CSV から分子グラフを構築
        # 実際は RDKit などで SMILES → グラフ変換を行う
        for i in range(100):  # 100個のダミー分子
            num_atoms = torch.randint(5, 20, (1,)).item()
            num_bonds = torch.randint(num_atoms - 1, num_atoms * 2, (1,)).item()

            # ノード特徴量（原子タイプ・電荷など）
            x = torch.randn(num_atoms, 9)
            # 辺（化学結合）
            src = torch.randint(0, num_atoms, (num_bonds,))
            dst = torch.randint(0, num_atoms, (num_bonds,))
            edge_index = torch.stack([src, dst], dim=0)
            # グラフラベル（毒性：0 or 1）
            y = torch.randint(0, 2, (1,))

            data = Data(x=x, edge_index=edge_index, y=y)
            data_list.append(data)

        # pre_transform を適用
        if self.pre_transform is not None:
            data_list = [self.pre_transform(d) for d in data_list]

        # 一括保存
        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

# 使用例（構文確認済み）
import tempfile
with tempfile.TemporaryDirectory() as tmpdir:
    my_dataset = MyMoleculeDataset(root=tmpdir)
    print(f"グラフ数: {len(my_dataset)}")           # 100
    print(f"最初のグラフ: {my_dataset[0]}")          # Data オブジェクト
    print(f"クラス数: {my_dataset.num_classes}")     # 2

from torch_geometric.data import Dataset

class LargeGraphDataset(Dataset):
    """大規模データセット: ファイルごとに遅延読み込み"""

    def __init__(self, root, file_list, transform=None):
        self.file_list = file_list
        super().__init__(root, transform)

    def len(self):
        return len(self.file_list)

    def get(self, idx):
        """インデックス idx のグラフを読み込んで返す"""
        # 各グラフをファイルから個別に読み込む
        data = torch.load(os.path.join(self.root, self.file_list[idx]))
        return data

# 対処1: バッチサイズを減らす
loader = DataLoader(dataset, batch_size=8)  # 32 → 8 に減らす

# 対処2: グラフをミニバッチに分割して処理（NeighborLoader）
from torch_geometric.loader import NeighborLoader
loader = NeighborLoader(
    data,
    num_neighbors=[25, 10],  # 各層で何近傍をサンプリングするか
    batch_size=1024,
    input_nodes=data.train_mask,
)

# 誤: edge_index が [E, 2] の形状
edge_index_wrong = torch.tensor([[0, 1], [1, 2], [2, 3]])
# print(edge_index_wrong.shape)  # [3, 2] ← NG

# 正: [2, E] の形状にする
edge_index_correct = edge_index_wrong.T  # 転置
# または
edge_index_correct = torch.tensor([[0, 1, 2], [1, 2, 3]])
# print(edge_index_correct.shape)  # [2, 3] ← OK

from torch_geometric.utils import to_undirected, is_undirected

# 有向グラフを無向グラフに変換（逆方向の辺を追加）
edge_index_undirected = to_undirected(edge_index)
print(is_undirected(edge_index_undirected))  # True

from torch_geometric.utils import add_self_loops, remove_self_loops

# GCNConv などで自己情報を集約するために必要
edge_index_with_loops, _ = add_self_loops(edge_index, num_nodes=data.num_nodes)

# 不要な自己ループを削除
edge_index_clean, _ = remove_self_loops(edge_index)

# 特徴量がない場合は次数または定数ベクトルで代用
if data.x is None:
    # 方法1: 全ノードを同じ定数特徴で初期化
    data.x = torch.ones(data.num_nodes, 1)
    # 方法2: 次数を特徴量として使う
    from torch_geometric.utils import degree
    deg = degree(data.edge_index[0], data.num_nodes, dtype=torch.float)
    data.x = deg.unsqueeze(-1)

# 訓練時
model.train()
# 評価時（Dropout が無効になる）
model.eval()
with torch.no_grad():
    pred = model(data.x, data.edge_index)

from torch_geometric.datasets import (
    Planetoid,    # Cora, Citeseer, PubMed (引用ネットワーク)
    TUDataset,    # MUTAG, ENZYMES, PROTEINS (グラフ分類)
    Amazon,       # Computer, Photo (商品共起グラフ)
    Coauthor,     # CS, Physics (共著ネットワーク)
    KarateClub,   # ザラテクラブ (チュートリアル定番)
    OGB,          # Open Graph Benchmark (大規模)
)

# 使用例
from torch_geometric import datasets
dataset = datasets.KarateClub()
data = dataset[0]
print(data.num_nodes, data.num_edges)  # 34, 156

pip install ogb

from ogb.nodeproppred import PygNodePropPredDataset
# ogbn-arxiv: 17万ノード・120万辺の引用グラフ
dataset = PygNodePropPredDataset("ogbn-arxiv", root="/tmp/ogb")

from torch_geometric.data import HeteroData

# 推薦グラフの例: ユーザー・映画・俳優の3種類ノード
hetero_data = HeteroData()

# 各ノードタイプの特徴量
hetero_data["user"].x = torch.randn(1000, 64)     # 1000 ユーザー
hetero_data["movie"].x = torch.randn(5000, 128)   # 5000 映画
hetero_data["actor"].x = torch.randn(10000, 32)   # 10000 俳優

# 辺タイプ: (送信ノードタイプ, 辺の種類, 受信ノードタイプ)
hetero_data["user", "watched", "movie"].edge_index = torch.randint(0, 5000, (2, 50000))
hetero_data["user", "rated", "movie"].edge_index = torch.randint(0, 5000, (2, 20000))
hetero_data["actor", "appeared_in", "movie"].edge_index = torch.randint(0, 5000, (2, 100000))

print(hetero_data)
# HeteroData(
#   user={ x=[1000, 64] },
#   movie={ x=[5000, 128] },
#   actor={ x=[10000, 32] },
#   (user, watched, movie)={ edge_index=[2, 50000] },
#   ...
# )
print(hetero_data.node_types)   # ['user', 'movie', 'actor']
print(hetero_data.edge_types)   # [('user', 'watched', 'movie'), ...]

from torch_geometric.nn import HeteroConv, SAGEConv

# HeteroConv: 辺タイプごとに異なる Conv を適用
conv = HeteroConv({
    ("user", "watched", "movie"): SAGEConv(64, 64),
    ("actor", "appeared_in", "movie"): SAGEConv(32, 64),
}, aggr="sum")  # 複数辺タイプからの集約を和で結合

# 辺タイプごとに異なる変換が適用される
# out_dict = conv(x_dict, edge_index_dict)
# out_dict["movie"]: 映画ノードの更新された特徴量

from torch_geometric.utils import negative_sampling

# 正例: 実際に存在する辺
pos_edge_index = data.edge_index  # [2, E]

# 負例: ランダムにサンプリングした存在しない辺
neg_edge_index = negative_sampling(
    edge_index=data.edge_index,
    num_nodes=data.num_nodes,
    num_neg_samples=data.num_edges,  # 正例と同数の負例
)

def link_pred_loss(z, pos_edge_index, neg_edge_index):
    """ノード埋め込み z からリンク予測損失を計算"""
    # 正例のスコア: ノードiとノードjの埋め込みの内積
    pos_score = (z[pos_edge_index[0]] * z[pos_edge_index[1]]).sum(dim=-1)
    neg_score = (z[neg_edge_index[0]] * z[neg_edge_index[1]]).sum(dim=-1)

    # BCE 損失
    pos_loss = F.binary_cross_entropy_with_logits(pos_score, torch.ones_like(pos_score))
    neg_loss = F.binary_cross_entropy_with_logits(neg_score, torch.zeros_like(neg_score))
    return pos_loss + neg_loss

# GNN でノード埋め込みを取得し、内積でリンク存在確率を推定
# z = gnn_model(data.x, train_edge_index)  # 訓練辺のみ使う
# loss = link_pred_loss(z, pos_edge_index, neg_edge_index)