图神经网络 GNN:GCN, GAT, GraphSAGE

import torch
import torch.nn as nn
from typing import Optional

class MessagePassingLayer(nn.Module):
    """通用消息传递层——GNN 的统一框架"""
    def __init__(self, in_dim, out_dim, edge_dim=0):
        super().__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.use_edge_attr = edge_dim > 0
        
        if self.use_edge_attr:
            self.msg_net = nn.Sequential(
                nn.Linear(in_dim * 2 + edge_dim, out_dim),
                nn.ReLU(),
                nn.Linear(out_dim, out_dim),
            )
        else:
            self.msg_net = nn.Sequential(
                nn.Linear(in_dim * 2, out_dim),
                nn.ReLU(),
                nn.Linear(out_dim, out_dim),
            )
        
        self.update_net = nn.Sequential(
            nn.Linear(in_dim + out_dim, out_dim),
            nn.ReLU(),
            nn.Linear(out_dim, out_dim),
        )
    
    def message(self, h_src, h_dst, edge_attr=None):
        """生成每条边的消息 m_uv"""
        if self.use_edge_attr and edge_attr is not None:
            return self.msg_net(torch.cat([h_src, h_dst, edge_attr], dim=-1))
        return self.msg_net(torch.cat([h_src, h_dst], dim=-1))
    
    def aggregate(self, messages, dst_idx, num_nodes):
        """聚合消息：mean pooling"""
        # dst_idx: 每条边目标节点的索引
        agg = torch.zeros(num_nodes, self.out_dim, device=messages.device)
        count = torch.zeros(num_nodes, 1, device=messages.device)
        agg.index_add_(0, dst_idx, messages)
        count.index_add_(0, dst_idx, torch.ones_like(messages[:, :1]))
        count = count.clamp(min=1)
        return agg / count
    
    def update(self, h_old, aggregated):
        """更新节点表示"""
        return self.update_net(torch.cat([h_old, aggregated], dim=-1))
    
    def forward(self, h, edge_index, edge_attr=None):
        """
        h: (N, in_dim) 节点特征
        edge_index: (2, E) [src; dst]
        edge_attr: (E, edge_dim) 可选边特征
        """
        src, dst = edge_index[0], edge_index[1]
        messages = self.message(h[src], h[dst], edge_attr)
        aggregated = self.aggregate(messages, dst, h.size(0))
        return self.update(h, aggregated)

# 构建一个简单图并测试消息传递
def build_sample_graph():
    """构建一个 6 节点的社交网络示例图"""
    # 节点特征: [活跃度, 粉丝数(归一化), 发帖频率]
    node_features = torch.tensor([
        [0.8, 0.9, 0.7],   # 用户 A: 高活跃大 V
        [0.3, 0.2, 0.5],   # 用户 B: 低活跃普通用户
        [0.9, 0.8, 0.9],   # 用户 C: 高活跃创作者
        [0.1, 0.1, 0.1],   # 用户 D: 潜水用户
        [0.6, 0.5, 0.8],   # 用户 E: 中等活跃
        [0.4, 0.3, 0.2],   # 用户 F: 低活跃浏览者
    ], dtype=torch.float32)
    
    # 边: (A->B, A->C, B->D, C->E, C->A, E->F, E->C, F->B)
    edge_index = torch.tensor([
        [0, 0, 1, 2, 2, 4, 4, 5],  # source
        [1, 2, 3, 4, 0, 5, 2, 1],  # target
    ], dtype=torch.long)
    
    return node_features, edge_index

# 测试
features, edges = build_sample_graph()
mp_layer = MessagePassingLayer(in_dim=3, out_dim=8)
output = mp_layer(features, edges)
print(f"输入特征: {features.shape} → 输出特征: {output.shape}")
print(f"\n消息传递后的节点表示（前 3 个节点）:")
for i in range(3):
    print(f"  节点 {chr(65+i)}: {output[i].detach().numpy()[:4].round(3)}...")

import torch
import torch.nn as nn
import torch.nn.functional as F

class GCNLayer(nn.Module):
    """单层 GCN：H' = σ(D^(-1/2) Ã D^(-1/2) H W)"""
    def __init__(self, in_dim, out_dim, dropout=0.5):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim, bias=False)
        self.dropout = nn.Dropout(dropout)
        self.reset_parameters()
    
    def reset_parameters(self):
        # Xavier 初始化
        nn.init.xavier_uniform_(self.linear.weight)
    
    def forward(self, h, adj_norm):
        """
        h: (N, in_dim) 节点特征
        adj_norm: (N, N) 归一化邻接矩阵 D^(-1/2) Ã D^(-1/2)
        """
        # 邻接矩阵聚合 + 线性变换
        support = self.linear(h)         # (N, out_dim)
        output = torch.spmm(adj_norm, support)  # (N, out_dim)
        return output


def compute_normalized_adjacency(adj):
    """计算 GCN 归一化邻接矩阵 D^(-1/2) Ã D^(-1/2)"""
    # 添加自环
    N = adj.size(0)
    A_tilde = adj + torch.eye(N, device=adj.device)
    
    # 计算度矩阵的 -1/2 次方
    degree = A_tilde.sum(dim=1)  # 每个节点的度
    degree_inv_sqrt = degree.pow(-0.5)
    degree_inv_sqrt[degree_inv_sqrt == float('inf')] = 0
    
    # D^(-1/2) * A~ * D^(-1/2)
    D_inv_sqrt = torch.diag(degree_inv_sqrt)
    adj_norm = D_inv_sqrt @ A_tilde @ D_inv_sqrt
    return adj_norm

class GCN(nn.Module):
    """两层 GCN 半监督节点分类器"""
    def __init__(self, in_dim, hidden_dim, out_dim, dropout=0.5):
        super().__init__()
        self.gc1 = GCNLayer(in_dim, hidden_dim, dropout)
        self.gc2 = GCNLayer(hidden_dim, out_dim, dropout)
    
    def forward(self, h, adj_norm):
        h = F.relu(self.gc1(h, adj_norm))
        h = F.dropout(h, p=0.5, training=self.training)
        h = self.gc2(h, adj_norm)
        return F.log_softmax(h, dim=1)


# === Cora 数据集模拟实验 ===
def simulate_cora_experiment():
    torch.manual_seed(42)
    N_NODES = 2708    # Cora 论文数
    N_FEATURES = 1433  # 词袋特征维度
    N_CLASSES = 7      # 主题类别
    N_TRAIN = 140      # 每类 20 个标签
    HIDDEN = 64
    
    # 模拟特征和图结构
    features = torch.randn(N_NODES, N_FEATURES)
    # 生成稀疏邻接矩阵（模拟 citation network）
    row_idx = torch.randint(0, N_NODES, (10000,))
    col_idx = torch.randint(0, N_NODES, (10000,))
    vals = torch.ones(10000)
    adj = torch.sparse.FloatTensor(
        torch.stack([row_idx, col_idx]), vals, (N_NODES, N_NODES))
    adj_dense = adj.to_dense()
    
    # 归一化
    adj_norm = compute_normalized_adjacency(adj_dense)
    
    # 构建模型
    model = GCN(N_FEATURES, HIDDEN, N_CLASSES)
    
    # 模拟训练标签
    labels = torch.randint(0, N_CLASSES, (N_NODES,))
    train_mask = torch.zeros(N_NODES, dtype=torch.bool)
    train_mask[:N_TRAIN] = True
    
    # 训练循环
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
    for epoch in range(200):
        model.train()
        optimizer.zero_grad()
        out = model(features, adj_norm)
        loss = F.nll_loss(out[train_mask], labels[train_mask])
        loss.backward()
        optimizer.step()
        
        if epoch % 50 == 0:
            model.eval()
            with torch.no_grad():
                preds = out.argmax(dim=1)
                acc = (preds[train_mask] == labels[train_mask]).float().mean()
            print(f"Epoch {epoch:3d}: Loss={loss:.4f}, Train Acc={acc:.3f}")

simulate_cora_experiment()

class GATLayer(nn.Module):
    """单层 GAT：注意力加权的邻居聚合"""
    def __init__(self, in_dim, out_dim, n_heads=1, dropout=0.6, alpha=0.2):
        super().__init__()
        self.n_heads = n_heads
        self.out_dim = out_dim
        self.dropout = nn.Dropout(dropout)
        self.leaky_relu = nn.LeakyReLU(alpha)
        
        # 每个头的线性变换和注意力向量
        self.W = nn.Parameter(torch.empty(n_heads, in_dim, out_dim))
        self.a_src = nn.Parameter(torch.empty(n_heads, out_dim, 1))
        self.a_dst = nn.Parameter(torch.empty(n_heads, out_dim, 1))
        
        self.reset_parameters()
    
    def reset_parameters(self):
        nn.init.xavier_uniform_(self.W)
        nn.init.xavier_uniform_(self.a_src)
        nn.init.xavier_uniform_(self.a_dst)
    
    def forward(self, h, edge_index):
        """
        h: (N, in_dim) 节点特征
        edge_index: (2, E) 边列表 [src; dst]
        """
        src, dst = edge_index[0], edge_index[1]
        N = h.size(0)
        
        # 线性变换: (n_heads, N, out_dim)
        Wh = torch.einsum('hid,nhd->nhid', h, self.W)
        
        # 计算注意力系数: 对每条边 (src->dst)
        # e_ij = aᵀ[Wh_i || Wh_j]
        e_src = (Wh[:, src] * self.a_src).sum(dim=-1)  # (n_heads, E, 1)
        e_dst = (Wh[:, dst] * self.a_dst).sum(dim=-1)  # (n_heads, E, 1)
        e = self.leaky_relu(e_src + e_dst).squeeze(-1)  # (n_heads, E)
        
        # softmax 归一化（在每个目标节点的入边上）
        e_max = e.max(dim=1, keepdim=True)[0]
        e = e - e_max
        e_exp = e.exp()
        
        # 按目标节点聚合分母
        denom = torch.zeros((self.n_heads, N), device=h.device)
        denom.index_add_(1, dst, e_exp)
        alpha = e_exp / (denom[:, dst] + 1e-10)  # (n_heads, E)
        alpha = self.dropout(alpha)
        
        # 加权聚合
        out = torch.zeros((self.n_heads, N, self.out_dim), device=h.device)
        for head in range(self.n_heads):
            msg = alpha[head:head+1, :].unsqueeze(-1) * Wh[head:head+1, dst]
            out[head].index_add_(0, dst, msg.squeeze(0))
        
        return out  # (n_heads, N, out_dim)

class GAT(nn.Module):
    """完整 GAT 模型：多头注意力 + 分类输出"""
    def __init__(self, in_dim, hidden_dim, out_dim, n_heads=8, 
                 hidden_heads=1, dropout=0.6):
        super().__init__()
        # 第一层：多头注意力，输出拼接
        self.gat1 = GATLayer(in_dim, hidden_dim, n_heads=n_heads, dropout=dropout)
        self.concat_dim = n_heads * hidden_dim
        self.dropout = nn.Dropout(dropout)
        
        # 第二层：单头注意力（或较少头），输出平均
        self.gat2 = GATLayer(self.concat_dim, out_dim, n_heads=hidden_heads, dropout=dropout)
    
    def forward(self, h, edge_index):
        # 第一层：多头 → 拼接 → ELU
        h1 = self.gat1(h, edge_index)  # (n_heads, N, hidden_dim)
        h1 = h1.view(h1.size(1), -1)   # (N, n_heads * hidden_dim)
        h1 = F.elu(h1)
        h1 = self.dropout(h1)
        
        # 第二层：单头 → 直接输出
        h2 = self.gat2(h1, edge_index)  # (1, N, out_dim)
        return h2.squeeze(0)  # (N, out_dim)


# 可视化注意力权重
def visualize_gat_attention(gat_model, h, edge_index, node_labels, ax=None):
    """提取并可视化 GAT 第一层的注意力权重"""
    with torch.no_grad():
        src, dst = edge_index[0], edge_index[1]
        Wh = torch.einsum('hid,nhd->nhid', h, gat_model.gat1.W)
        
        e_src = (Wh[:, src] * gat_model.gat1.a_src).sum(dim=-1)
        e_dst = (Wh[:, dst] * gat_model.gat1.a_dst).sum(dim=-1)
        e = F.leaky_relu(e_src + e_dst, negative_slope=0.2).squeeze(-1)
        alpha = F.softmax(e, dim=1)  # (n_heads, E)
    
    print(f"注意力权重形状: {alpha.shape}")
    print(f"  n_heads={alpha.size(0)}, n_edges={alpha.size(1)}")
    
    # 每个头的平均注意力分布
    for head in range(alpha.size(0)):
        head_attn = alpha[head].detach().cpu().numpy()
        print(f"  Head {head}: mean={head_attn.mean():.4f}, "
              f"max={head_attn.max():.4f}, min={head_attn.min():.4f}, "
              f"std={head_attn.std():.4f}")
    
    return alpha.detach()


# 模拟实验
print("GAT 注意力权重分析:")
torch.manual_seed(42)
h = torch.randn(100, 16)
edges = torch.randint(0, 100, (2, 300))
gat = GAT(in_dim=16, hidden_dim=8, out_dim=7, n_heads=4)
attn_weights = visualize_gat_attention(gat, h, edges, None)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List

class GraphSAGEAggregator(nn.Module):
    """GraphSAGE 聚合器：支持 Mean / Pooling / LSTM"""
    def __init__(self, in_dim, out_dim, aggregator_type="mean"):
        super().__init__()
        self.aggregator_type = aggregator_type
        self.out_dim = out_dim
        
        if aggregator_type == "mean":
            pass  # 无需额外参数
        elif aggregator_type == "pooling":
            self.fc = nn.Linear(in_dim, in_dim)
        elif aggregator_type == "lstm":
            self.lstm = nn.LSTM(in_dim, in_dim, batch_first=True)
        else:
            raise ValueError(f"Unknown aggregator: {aggregator_type}")
        
        # GraphSAGE 的 W · concat(self, agg(neighbors))
        self.weight = nn.Linear(2 * in_dim, out_dim)
    
    def forward(self, self_features, neighbor_features):
        """
        self_features: (N, in_dim) 中心节点特征
        neighbor_features: (N, num_neighbors, in_dim) 邻居特征
        """
        if self.aggregator_type == "mean":
            # 简单平均
            aggregated = neighbor_features.mean(dim=1)  # (N, in_dim)
        elif self.aggregator_type == "pooling":
            # Pooling: 先变换再 max pooling
            transformed = F.relu(self.fc(neighbor_features))  # (N, K, in_dim)
            aggregated, _ = transformed.max(dim=1)  # (N, in_dim)
        elif self.aggregator_type == "lstm":
            # LSTM: 邻居作为序列
            lstm_out, _ = self.lstm(neighbor_features)
            aggregated = lstm_out[:, -1, :]  # 取最后一个时间步
        
        # 拼接自身和聚合的邻居特征
        combined = torch.cat([self_features, aggregated], dim=1)
        return F.normalize(self.weight(combined), dim=1)  # L2 归一化


class NeighborSampler:
    """邻居采样器"""
    def __init__(self, adj_dict):
        self.adj_dict = adj_dict  # {node_id: [neighbor_ids]}
    
    def sample_neighbors(self, node_ids, sample_size):
        """为每个节点采样固定数量的邻居"""
        sampled = []
        for node in node_ids:
            neighbors = self.adj_dict.get(node, [])
            if len(neighbors) <= sample_size:
                sampled.append(neighbors)
            else:
                import random
                sampled.append(random.sample(neighbors, sample_size))
        return sampled
    
    def get_batch_features(self, node_ids, neighbor_lists, features):
        """收集批次中所有节点和邻居的特征"""
        self_feat = features[node_ids]
        neighbor_feat_list = []
        for neighbors in neighbor_lists:
            if len(neighbors) == 0:
                neighbor_feat_list.append(torch.zeros(1, features.size(1)))
            else:
                neighbor_feat_list.append(features[neighbors])
        return self_feat, neighbor_feat_list

class GraphSAGE(nn.Module):
    """两层 GraphSAGE 模型"""
    def __init__(self, in_dim, hidden_dim, out_dim, aggregator_type="mean"):
        super().__init__()
        self.agg1 = GraphSAGEAggregator(in_dim, hidden_dim, aggregator_type)
        self.agg2 = GraphSAGEAggregator(hidden_dim, out_dim, aggregator_type)
    
    def forward(self, h1, neighbors1, h2, neighbors2):
        """
        h1: 中心节点特征, neighbors1: 一跳邻居特征
        h2: 一跳节点特征, neighbors2: 二跳邻居特征
        """
        # 第一层：聚合二跳邻居 → 生成一跳节点表示
        h2_agg = self.agg1(h2, neighbors2)
        
        # 第二层：聚合一跳邻居（使用更新后的一跳表示） → 生成中心节点表示
        h1_agg = self.agg2(h1, h2_agg)
        return h1_agg
    
    @torch.no_grad()
    def inference(self, features, edge_index, batch_size=256):
        """全图推理（不采样，使用全部邻居）"""
        N = features.size(0)
        adj_dict = {}
        src, dst = edge_index
        for i in range(src.size(0)):
            adj_dict.setdefault(src[i].item(), []).append(dst[i].item())
        
        sampler = NeighborSampler(adj_dict)
        max_neighbors = max(len(v) for v in adj_dict.values()) if adj_dict else 0
        
        results = []
        for start in range(0, N, batch_size):
            end = min(start + batch_size, N)
            nodes = torch.arange(start, end)
            neighbor_lists = sampler.sample_neighbors(
                nodes.tolist(), max_neighbors)
            self_feat, neigh_feat = sampler.get_batch_features(
                nodes, neighbor_lists, features)
            
            # 确保邻居特征形状一致
            max_n = max(len(n) for n in neighbor_lists)
            padded = []
            for nl in neighbor_lists:
                nf = features[nl] if len(nl) > 0 else torch.zeros(1, features.size(1))
                if nf.size(0) < max_n:
                    pad = torch.zeros(max_n - nf.size(0), nf.size(1))
                    nf = torch.cat([nf, pad], dim=0)
                padded.append(nf)
            neigh_feat = torch.stack(padded)
            
            out = self.agg1(self_feat, neigh_feat)
            results.append(out)
        
        return torch.cat(results, dim=0)


# 测试 GraphSAGE
torch.manual_seed(42)
N = 500
in_dim = 64
features = torch.randn(N, in_dim)
edge_index = torch.stack([
    torch.randint(0, N, (2000,)),
    torch.randint(0, N, (2000,)),
])

model = GraphSAGE(in_dim, 32, 16)
print("GraphSAGE 模型初始化完成")
print(f"  聚合器类型: mean")
print(f"  参数量: {sum(p.numel() for p in model.parameters()):,}")

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleReadout(nn.Module):
    """基础图级读出：支持 mean / sum / max"""
    def __init__(self, mode="mean"):
        super().__init__()
        self.mode = mode
    
    def forward(self, h, batch_idx):
        """
        h: (N, d) 所有节点特征
        batch_idx: (N,) 每个节点所属的图 ID
        返回: (B, d) 图级表示
        """
        if self.mode == "mean":
            out = scatter_mean(h, batch_idx, dim=0)
        elif self.mode == "sum":
            out = scatter_sum(h, batch_idx, dim=0)
        elif self.mode == "max":
            out = scatter_max(h, batch_idx, dim=0)[0]
        return out


class SAGPoolLayer(nn.Module):
    """自注意力图池化（SAGPool）"""
    def __init__(self, in_dim, ratio=0.5):
        super().__init__()
        self.ratio = ratio
        self.score_net = nn.Sequential(
            nn.Linear(in_dim, in_dim),
            nn.Tanh(),
            nn.Linear(in_dim, 1),
        )
    
    def forward(self, h, edge_index, batch_idx):
        """
        h: (N, in_dim), edge_index: (2, E)
        返回: 池化后的节点特征、边索引、batch 索引
        """
        scores = self.score_net(h).squeeze(-1)  # (N,)
        k = max(1, int(scores.size(0) * self.ratio))
        
        # top-k 选择
        top_idx = scores.topk(k, dim=0)[1]
        top_idx = top_idx.sort()[0]
        
        # 选择保留的节点
        h_pool = h[top_idx]
        batch_idx_pool = batch_idx[top_idx]
        
        # 过滤边：只保留两端节点都在 top-k 中的边
        src, dst = edge_index
        mask = torch.isin(src, top_idx) & torch.isin(dst, top_idx)
        edge_index_pool = edge_index[:, mask]
        
        # 注意力加权
        attention_scores = torch.sigmoid(scores[top_idx])
        h_pool = h_pool * attention_scores.unsqueeze(-1)
        
        return h_pool, edge_index_pool, batch_idx_pool


def scatter_mean(src, index, dim=0):
    """按索引求均值（简化版 scatter_mean）"""
    out = torch.zeros((index.max().item() + 1, src.size(1)), 
                      device=src.device, dtype=src.dtype)
    count = torch.zeros(out.size(0), device=src.device)
    out.index_add_(0, index, src)
    count.index_add_(0, index, torch.ones(src.size(0), device=src.device))
    return out / count.clamp(min=1).unsqueeze(-1)


def scatter_sum(src, index, dim=0):
    """按索引求和（简化版 scatter_sum）"""
    out = torch.zeros((index.max().item() + 1, src.size(1)),
                      device=src.device, dtype=src.dtype)
    out.index_add_(0, index, src)
    return out

class GraphClassifier(nn.Module):
    """图分类模型：GCN + 多层池化 + 分类头"""
    def __init__(self, in_dim, hidden_dim, out_dim, pool_ratios=[0.8, 0.5]):
        super().__init__()
        self.gc1 = GCNLayer(in_dim, hidden_dim)
        self.gc2 = GCNLayer(hidden_dim, hidden_dim)
        
        self.pool1 = SAGPoolLayer(hidden_dim, ratio=pool_ratios[0])
        self.pool2 = SAGPoolLayer(hidden_dim, ratio=pool_ratios[1])
        
        self.readout = SimpleReadout(mode="mean")
        
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(hidden_dim, out_dim),
        )
    
    def forward(self, h, edge_index, batch_idx):
        # GCN 编码
        h = F.relu(self.gc1(h, adj_norm=None))
        h = F.relu(self.gc2(h, adj_norm=None))
        
        # 第一层池化
        h1_pool, edge1, batch1 = self.pool1(h, edge_index, batch_idx)
        
        # 第二层池化
        h2_pool, edge2, batch2 = self.pool2(h1_pool, edge1, batch1)
        
        # 多尺度读出：池化前 + 池化后
        h_global = self.readout(h, batch_idx)
        h_fine = self.readout(h2_pool, batch2)
        
        h_graph = torch.cat([h_global, h_fine], dim=1)
        return self.classifier(h_graph)


# 测试图分类
def test_graph_classification():
    torch.manual_seed(42)
    
    # 模拟 32 个图，每个图 10-50 个节点
    graphs = []
    for i in range(32):
        n = torch.randint(10, 50, (1,)).item()
        h = torch.randn(n, 16)
        edges = torch.randint(0, n, (2, n * 3))
        graphs.append((h, edges))
    
    # 合并为一个批次
    all_h = torch.cat([g[0] for g in graphs], dim=0)
    all_edges = torch.cat([g[1] + torch.tensor([[sum(gg[0].size(0) for gg in graphs[:i])],
                                                   [sum(gg[0].size(0) for gg in graphs[:i])]])
                            for i, g in enumerate(graphs)], dim=1)
    batch_idx = torch.cat([torch.full((g[0].size(0),), i) 
                           for i, g in enumerate(graphs)])
    
    model = GraphClassifier(in_dim=16, hidden_dim=32, out_dim=5)
    out = model(all_h, all_edges, batch_idx)
    print(f"图分类输出: {out.shape} (期望: [32, 5])")
    print(f"参数总量: {sum(p.numel() for p in model.parameters()):,}")

test_graph_classification()

# === 场景 1: GNN 推荐系统（简化版 PinSage）===
import torch
import torch.nn as nn

class GNNRecommender(nn.Module):
    """基于 GNN 的推荐系统"""
    def __init__(self, n_users, n_items, embed_dim, n_layers=2):
        super().__init__()
        self.user_embedding = nn.Embedding(n_users, embed_dim)
        self.item_embedding = nn.Embedding(n_items, embed_dim)
        
        self.gnn_layers = nn.ModuleList([
            GCNLayer(embed_dim, embed_dim) for _ in range(n_layers)
        ])
        self.dropout = nn.Dropout(0.3)
    
    def build_bipartite_adj(self, interactions, n_users, n_items):
        """构建用户-商品二分图邻接矩阵"""
        N = n_users + n_items
        adj = torch.zeros(N, N)
        for u, i in interactions:
            adj[u, n_users + i] = 1.0
            adj[n_users + i, u] = 1.0
        return compute_normalized_adjacency(adj)
    
    def forward(self, user_ids, item_ids, adj_norm):
        # 初始化嵌入
        all_users = self.user_embedding.weight
        all_items = self.item_embedding.weight
        all_embed = torch.cat([all_users, all_items], dim=0)
        
        # GNN 传播
        h = all_embed
        all_h = [h]
        for gnn in self.gnn_layers:
            h = F.relu(gnn(h, adj_norm))
            h = self.dropout(h)
            all_h.append(h)
        
        # 多层聚合（LightGCN 风格：各层平均）
        final_h = sum(all_h) / len(all_h)
        
        # 查询特定用户和商品的表示
        u_emb = final_h[user_ids]
        i_emb = final_h[n_users + item_ids]
        return u_emb, i_emb
    
    def predict_score(self, user_ids, item_ids, adj_norm):
        u_emb, i_emb = self.forward(user_ids, item_ids, adj_norm)
        return (u_emb * i_emb).sum(dim=1)  # 内积打分

# === 场景 2: 分子性质预测 ===
class MoleculeGNN(nn.Module):
    """分子图神经网络：原子类型 + 键类型编码"""
    def __init__(self, n_atom_types=100, n_bond_types=6, hidden_dim=64, out_dim=1):
        super().__init__()
        self.atom_embedding = nn.Embedding(n_atom_types, hidden_dim)
        self.bond_embedding = nn.Embedding(n_bond_types, hidden_dim)
        
        self.gnn_layers = nn.ModuleList([
            nn.ModuleDict({
                "node_update": nn.Linear(hidden_dim * 2, hidden_dim),
                "edge_update": nn.Linear(hidden_dim * 3, hidden_dim),
            })
            for _ in range(3)
        ])
        
        self.readout = SimpleReadout(mode="sum")
        self.predictor = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim // 2, out_dim),
        )
    
    def forward(self, atom_types, edge_index, bond_types, batch_idx):
        # 嵌入
        h_atom = self.atom_embedding(atom_types)
        h_bond = self.bond_embedding(bond_types)
        
        # GNN 消息传递（带边特征）
        for layer in self.gnn_layers:
            src, dst = edge_index
            
            # 消息: h_u + h_v + e_uv
            messages = h_atom[src] + h_atom[dst] + h_bond
            
            # 聚合
            agg = scatter_mean(messages, dst, dim=0)
            
            # 更新
            h_atom = layer["node_update"](
                torch.cat([h_atom, agg], dim=-1))
            h_atom = F.relu(h_atom)
        
        # 图级读出 + 预测
        h_graph = self.readout(h_atom, batch_idx)
        return self.predictor(h_graph).squeeze(-1)


# === 场景 3: 社交网络异常检测 ===
def detect_fraud_users(node_features, edge_index, n_anomalies=50):
    """基于图结构和节点特征的异常用户检测"""
    # 计算图结构特征
    src, dst = edge_index
    degree = torch.zeros(node_features.size(0))
    degree.index_add_(0, src, torch.ones(src.size(0)))
    
    # 聚类系数（局部密度）
    local_density = torch.zeros(node_features.size(0))
    for i in range(src.size(0)):
        local_density[src[i]] += 1.0 / (degree[src[i]] + 1)
    
    # 结合图特征和原始特征
    structural_features = torch.stack([degree, local_density], dim=1)
    combined = torch.cat([node_features, structural_features], dim=1)
    
    # GNN 编码
    adj = torch.zeros(node_features.size(0), node_features.size(0))
    adj[src, dst] = 1.0
    adj_norm = compute_normalized_adjacency(adj)
    
    model = GCNLayer(combined.size(1), 32)
    h = F.relu(model(combined, adj_norm))
    
    # 异常分数（基于重构误差或偏离度）
    anomaly_scores = torch.norm(h - h.mean(dim=0), dim=1)
    top_anomalies = anomaly_scores.topk(n_anomalies)[1]
    
    return top_anomalies, anomaly_scores


print("三大应用场景模型定义完成:")
print("  1. GNN 推荐系统 (PinSage 风格)")
print("  2. 分子性质预测 (带边特征的 GNN)")
print("  3. 社交网络异常检测 (结构特征 + GNN)")

# === PyG 实战 1: 节点分类 (GCN on Cora) ===
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeatures

def train_gcn_on_cora():
    """在 Cora 数据集上训练 GCN 节点分类器"""
    # 加载数据
    dataset = Planetoid(root='/tmp/Cora', name='Cora',
                        transform=NormalizeFeatures())
    data = dataset[0]  # Cora 是单个图
    print(f"Cora 数据集: {data.num_nodes} 节点, {data.num_edges} 边, "
          f"{data.num_features} 维特征, {dataset.num_classes} 类")
    
    class GCN(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, out_channels)
        
        def forward(self, x, edge_index):
            x = self.conv1(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, p=0.5, training=self.training)
            x = self.conv2(x, edge_index)
            return F.log_softmax(x, dim=1)
    
    model = GCN(data.num_features, 16, dataset.num_classes)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
    
    model.train()
    for epoch in range(200):
        optimizer.zero_grad()
        out = model(data.x, data.edge_index)
        loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
        loss.backward()
        optimizer.step()
    
    # 评估
    model.eval()
    with torch.no_grad():
        pred = model(data.x, data.edge_index).argmax(dim=1)
        test_correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
        test_total = data.test_mask.sum()
        test_acc = int(test_correct) / int(test_total)
    
    print(f"测试准确率: {test_acc:.4f}")
    return model

# 模拟运行（不依赖 PyG 安装）
print("GCN on Cora 代码示例已定义（需安装 torch_geometric 运行）")
print("  pip install torch-geometric")

# === PyG 实战 2: 图分类 (GIN on MUTAG) ===
from torch_geometric.nn import GINConv, global_add_pool
from torch_geometric.loader import DataLoader
from torch_geometric.datasets import TUDataset

def train_gin_on_mutag():
    """在 MUTAG 数据集上训练 GIN 图分类器"""
    # 加载数据
    dataset = TUDataset(root='/tmp/MUTAG', name='MUTAG')
    print(f"MUTAG 数据集: {len(dataset)} 个图, "
          f"类别数: {dataset.num_classes}")
    print(f"  平均节点数: {dataset.data.num_nodes / len(dataset):.1f}")
    
    # 划分训练/测试集
    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)
    
    class GIN(torch.nn.Module):
        """Graph Isomorphism Network (GIN)"""
        def __init__(self, in_dim, hidden_dim, out_dim, n_layers=5):
            super().__init__()
            self.n_layers = n_layers
            
            # GIN 层: h' = MLP((1+eps) * h + sum(h_neighbors))
            self.convs = torch.nn.ModuleList()
            self.batch_norms = torch.nn.ModuleList()
            
            for i in range(n_layers):
                in_c = in_dim if i == 0 else hidden_dim
                mlp = torch.nn.Sequential(
                    torch.nn.Linear(in_c, hidden_dim),
                    torch.nn.ReLU(),
                    torch.nn.Linear(hidden_dim, hidden_dim),
                )
                self.convs.append(GINConv(mlp, train_eps=True))
                self.batch_norms.append(torch.nn.BatchNorm1d(hidden_dim))
            
            self.pool = global_add_pool
            self.classifier = torch.nn.Sequential(
                torch.nn.Linear(hidden_dim, hidden_dim // 2),
                torch.nn.ReLU(),
                torch.nn.Dropout(0.5),
                torch.nn.Linear(hidden_dim // 2, out_dim),
            )
        
        def forward(self, x, edge_index, batch):
            h = x
            for i in range(self.n_layers):
                h = self.convs[i](h, edge_index)
                h = self.batch_norms[i](h)
                if i < self.n_layers - 1:
                    h = F.relu(h)
            
            h_graph = self.pool(h, batch)
            return self.classifier(h_graph)
    
    model = GIN(dataset.num_features, 32, dataset.num_classes)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
    # 训练循环
    for epoch in range(100):
        model.train()
        total_loss = 0
        for batch_data in train_loader:
            optimizer.zero_grad()
            out = model(batch_data.x, batch_data.edge_index, batch_data.batch)
            loss = F.cross_entropy(out, batch_data.y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * batch_data.num_graphs
        
        # 测试
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for batch_data in test_loader:
                out = model(batch_data.x, batch_data.edge_index, batch_data.batch)
                correct += (out.argmax(dim=1) == batch_data.y).sum().item()
                total += batch_data.num_graphs
        
        if epoch % 20 == 0:
            print(f"Epoch {epoch:3d}: Loss={total_loss/len(train_dataset):.4f}, "
                  f"Test Acc={correct/total:.3f}")
    
    print(f"最终测试准确率: {correct/total:.4f}")
    return model

print("GIN on MUTAG 代码示例已定义（需安装 torch_geometric 运行）")
print("  pip install torch-geometric")