DAG-GNN: Directed Acyclic Graph - Graph Neural Network¶

for epoch in range(max_epochs):
    # 前向传播
    H = gnn_encoder(X, W, theta)  # 编码
    X_hat = gnn_decoder(H, theta)  # 解码

    # 计算损失
    loss_recon = reconstruction_loss(X, X_hat)
    loss_dag = dag_constraint(W)
    loss_sparse = l1_norm(W)

    total_loss = loss_recon + lambda1 * loss_dag + lambda2 * loss_sparse

    # 反向传播
    total_loss.backward()
    optimizer.step()

    # 输出当前结果
    if epoch % log_interval == 0:
        print(f"Epoch {epoch}, Loss: {total_loss.item()}")

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

class DAG_GNN(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers):
        super(DAG_GNN, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers

        # GNN层
        self.gnn_layers = nn.ModuleList([
            nn.Linear(input_dim, hidden_dim) if i == 0 
            else nn.Linear(hidden_dim, hidden_dim)
            for i in range(num_layers)
        ])

        # 解码器
        self.decoder = nn.Linear(hidden_dim, input_dim)

        # 邻接矩阵（可学习参数）
        self.W = nn.Parameter(torch.randn(input_dim, input_dim))

    def forward(self, X):
        batch_size, num_nodes, feature_dim = X.shape

        # GNN编码
        H = X
        for layer in self.gnn_layers:
            # 图卷积操作
            H_new = torch.zeros_like(H)
            for i in range(num_nodes):
                # 聚合邻居信息
                neighbor_sum = torch.zeros(batch_size, self.hidden_dim)
                for j in range(num_nodes):
                    if i != j:
                        weight = torch.sigmoid(self.W[i, j])
                        neighbor_sum += weight * H[:, j, :]

                # 自身信息 + 邻居信息
                H_new[:, i, :] = layer(H[:, i, :] + neighbor_sum)

            H = torch.relu(H_new)

        # 解码
        X_hat = self.decoder(H)

        return X_hat, self.W

    def dag_constraint(self):
        """计算DAG约束"""
        W = self.W
        # 使用截断指数函数近似
        exp_W = torch.matrix_exp(W * W)
        return torch.trace(exp_W) - self.input_dim

def train_dag_gnn(data, hidden_dim=64, num_layers=3, 
                 lambda1=0.1, lambda2=0.01, lr=0.001, 
                 epochs=1000, batch_size=32):
    """
    训练DAG-GNN模型
    """
    input_dim = data.shape[1]
    model = DAG_GNN(input_dim, hidden_dim, num_layers)
    optimizer = optim.Adam(model.parameters(), lr=lr)

    # 数据准备
    dataset = prepare_dataset(data, batch_size)

    for epoch in range(epochs):
        total_loss = 0

        for batch_X in dataset:
            optimizer.zero_grad()

            # 前向传播
            X_hat, W = model(batch_X)

            # 计算损失
            recon_loss = nn.MSELoss()(X_hat, batch_X)
            dag_loss = model.dag_constraint()
            sparse_loss = torch.norm(W, p=1)

            total_loss_batch = recon_loss + lambda1 * dag_loss + lambda2 * sparse_loss

            # 反向传播
            total_loss_batch.backward()
            optimizer.step()

            total_loss += total_loss_batch.item()

        if epoch % 100 == 0:
            print(f"Epoch {epoch}, Loss: {total_loss/len(dataset):.4f}")

    return model

def prepare_dataset(data, batch_size):
    """准备数据集"""
    # 将数据重塑为 (batch_size, num_nodes, feature_dim)
    # 这里简化处理，实际需要根据具体数据结构调整
    n_samples = data.shape[0]
    n_vars = data.shape[1]

    # 为每个样本创建图结构
    dataset = []
    for i in range(0, n_samples - batch_size + 1, batch_size):
        batch_data = data[i:i+batch_size]
        # 重塑为 (batch_size, num_vars, 1)
        batch_reshaped = batch_data.unsqueeze(-1)
        dataset.append(batch_reshaped)

    return dataset

# 生成线性因果数据
n_samples = 1000
X1 = np.random.normal(0, 1, n_samples)
X2 = 0.8 * X1 + np.random.normal(0, 0.5, n_samples)
X3 = 0.6 * X2 + np.random.normal(0, 0.3, n_samples)

data = np.column_stack([X1, X2, X3])

# 训练DAG-GNN
model = train_dag_gnn(data, hidden_dim=32, num_layers=2, 
                     lambda1=0.1, lambda2=0.01, lr=0.001)

# 获取学习到的邻接矩阵
W_learned = model.W.detach().numpy()
print("Learned adjacency matrix:")
print(W_learned)

# 生成非线性因果数据
X1 = np.random.normal(0, 1, n_samples)
X2 = np.sin(X1) + np.random.normal(0, 0.3, n_samples)
X3 = X2**2 + np.random.normal(0, 0.2, n_samples)

data = np.column_stack([X1, X2, X3])