MindSpore 因果推理与因果发现实战

因果推理是人工智能领域的重要研究方向，它超越了传统相关性分析，致力于揭示变量之间的因果关系。本文将深入探讨如何使用 MindSpore 实现因果发现算法和因果推理模型，帮助读者掌握这一前沿技术。

一、因果推理基础理论

1.1 相关性与因果性的区别

在机器学习中，我们经常遇到"相关性不等于因果性"的问题。例如，冰淇淋销量与溺水事件呈正相关，但这并不意味着吃冰淇淋导致溺水——真正的共同原因是夏季高温。

因果推理的核心任务是回答"如果干预 X，Y 会如何变化"这类反事实问题，而不仅仅是预测观测数据中的关联模式。

1.2 结构因果模型（SCM）

结构因果模型是因果推理的数学基础，由以下要素组成：

• 外生变量 U：模型外部的随机变量
• 内生变量 V：由其他变量决定的变量
• 结构方程 f：描述变量间的函数关系
• 因果图 G：有向无环图（DAG），表示因果结构

# 因果图示例
#     U1      U2
#      ↓       ↓
#      X ───→ Y ───→ Z
#
# 表示：X 影响 Y，Y 影响 Z，X 和 Y 还受外部因素影响

1.3 干预与 do-演算

Judea Pearl 提出的 do-演算是因果推理的核心工具。do(X=x) 表示强制将 X 设为 x，这不同于观察到 X=x：

• 观察 P(Y|X=x)：在 X 自然取值为 x 的条件下 Y 的分布
• 干预 P(Y|do(X=x))：强制设置 X=x 后 Y 的分布

二、MindSpore 因果发现算法实现

2.1 PC 算法实现

PC 算法是经典的因果发现算法，通过条件独立性检验从观测数据中学习因果结构。

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Tensor
import numpy as np
from scipy import stats
from itertools import combinations

class PCSAlgorithm:
    """PC算法实现：从观测数据中发现因果结构"""

    def __init__(self, alpha=0.05):
        """
        初始化PC算法

        Args:
            alpha: 条件独立性检验的显著性水平
        """
        self.alpha = alpha
        self.graph = None  # 邻接矩阵表示的因果图
        self.sep_set = None  # 分离集

    def _conditional_independence_test(self, data, x, y, z_set):
        """
        条件独立性检验：X ⊥ Y | Z

        使用偏相关系数和Fisher z变换
        """
        n = data.shape[0]

        if len(z_set) == 0:
            # 无条件独立性检验：简单相关系数
            corr = np.corrcoef(data[:, x], data[:, y])[0, 1]
            z_value = 0.5 * np.log((1 + corr) / (1 - corr))
            se = 1.0 / np.sqrt(n - 3)
            p_value = 2 * (1 - stats.norm.cdf(abs(z_value) / se))
        else:
            # 条件独立性检验：偏相关系数
            z_indices = list(z_set)
            all_indices = [x, y] + z_indices
            sub_data = data[:, all_indices]

            # 计算协方差矩阵的逆（精度矩阵）
            cov = np.cov(sub_data.T)
            try:
                precision = np.linalg.inv(cov)
            except np.linalg.LinAlgError:
                return True, 1.0  # 奇异矩阵，认为独立

            # 偏相关系数 = -precision[i,j] / sqrt(precision[i,i] * precision[j,j])
            p_corr = -precision[0, 1] / np.sqrt(precision[0, 0] * precision[1, 1])

            # Fisher z变换
            z_value = 0.5 * np.log((1 + p_corr) / (1 - p_corr))
            dof = n - len(z_set) - 3
            se = 1.0 / np.sqrt(dof)
            p_value = 2 * (1 - stats.norm.cdf(abs(z_value) / se))

        return p_value > self.alpha, p_value

    def fit(self, data):
        """
        运行PC算法学习因果结构

        Args:
            data: 观测数据，shape=(n_samples, n_variables)
        """
        n_vars = data.shape[1]

        # Step 1: 初始化完全无向图
        self.graph = np.ones((n_vars, n_vars), dtype=int)
        np.fill_diagonal(self.graph, 0)

        # 分离集初始化
        self.sep_set = [[set() for _ in range(n_vars)] for _ in range(n_vars)]

        # Step 2: 边删除阶段
        depth = 0
        while True:
            removed_any = False

            for i in range(n_vars):
                for j in range(i + 1, n_vars):
                    if self.graph[i, j] == 0:
                        continue

                    # 获取i和j的邻接节点
                    adj_i = set(np.where(self.graph[i, :] == 1)[0]) - {j}
                    adj_j = set(np.where(self.graph[j, :] == 1)[0]) - {i}
                    adj = adj_i.union(adj_j)

                    if len(adj) < depth:
                        continue

                    # 检查所有大小为depth的子集
                    for subset in combinations(adj, depth):
                        is_indep, p_val = self._conditional_independence_test(
                            data, i, j, set(subset)
                        )
                        if is_indep:
                            # 删除边
                            self.graph[i, j] = 0
                            self.graph[j, i] = 0
                            # 记录分离集
                            self.sep_set[i][j] = set(subset)
                            self.sep_set[j][i] = set(subset)
                            removed_any = True
                            break

            if not removed_any:
                break
            depth += 1

        # Step 3: 方向确定阶段（v-structures）
        for i in range(n_vars):
            for j in range(n_vars):
                if self.graph[i, j] == 0:
                    continue
                for k in range(n_vars):
                    if k == i or k == j:
                        continue
                    if self.graph[j, k] == 0:
                        continue
                    if self.graph[i, k] == 1:
                        continue  # i-j-k 不构成路径

                    # 检查是否构成 v-structure: i -> j <- k
                    if i not in self.sep_set[i][k] and j not in self.sep_set[i][k]:
                        # 标记方向
                        self.graph[i, j] = -1  # i -> j
                        self.graph[j, i] = 1   # j <- i
                        self.graph[j, k] = -1  # j <- k
                        self.graph[k, j] = 1   # k -> j

        return self

    def get_causal_graph(self):
        """返回学习到的因果图"""
        return self.graph

    def print_edges(self, var_names=None):
        """打印因果边"""
        n = self.graph.shape[0]
        if var_names is None:
            var_names = [f'X{i}' for i in range(n)]

        edges = []
        for i in range(n):
            for j in range(n):
                if self.graph[i, j] == -1:  # i -> j
                    edges.append(f'{var_names[i]} -> {var_names[j]}')

        print("发现的因果边：")
        for edge in edges:
            print(f"  {edge}")

        # 未定向的边
        undirected = []
        for i in range(n):
            for j in range(i + 1, n):
                if self.graph[i, j] == 1 and self.graph[j, i] == 1:
                    undirected.append(f'{var_names[i]} - {var_names[j]}')

        if undirected:
            print("\n未定向的边：")
            for edge in undirected:
                print(f"  {edge}")


# 示例：使用PC算法发现因果关系
def demo_pc_algorithm():
    """演示PC算法在模拟数据上的因果发现"""
    np.random.seed(42)
    n_samples = 1000

    # 生成模拟因果数据
    # 真实因果结构：X1 -> X3, X2 -> X3, X3 -> X4
    U1 = np.random.randn(n_samples)
    U2 = np.random.randn(n_samples)
    U3 = np.random.randn(n_samples)
    U4 = np.random.randn(n_samples)

    X1 = U1
    X2 = U2
    X3 = 0.5 * X1 + 0.7 * X2 + U3
    X4 = 0.8 * X3 + U4

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

    # 运行PC算法
    pc = PCSAlgorithm(alpha=0.01)
    pc.fit(data)

    var_names = ['X1', 'X2', 'X3', 'X4']
    pc.print_edges(var_names)

    return pc


if __name__ == '__main__':
    demo_pc_algorithm()

2.2 GES 算法（贪婪等价搜索）

GES 算法通过优化评分函数在DAG空间中搜索最优因果结构。

import mindspore as ms
from mindspore import Tensor
import numpy as np

class GESAlgorithm:
    """贪婪等价搜索算法实现"""

    def __init__(self, score_type='bic'):
        """
        Args:
            score_type: 评分类型，'bic' 或 'bdeu'
        """
        self.score_type = score_type
        self.graph = None

    def _compute_local_score(self, data, node, parents):
        """计算局部BIC评分"""
        n = data.shape[0]
        y = data[:, node]

        if len(parents) == 0:
            # 无父节点：估计边缘分布的方差
            var = np.var(y)
            score = -n / 2 * np.log(2 * np.pi * var) - n / 2
        else:
            # 有父节点：线性回归的残差方差
            X = data[:, list(parents)]
            # 添加截距项
            X = np.column_stack([np.ones(n), X])

            # 最小二乘解
            try:
                beta = np.linalg.lstsq(X, y, rcond=None)[0]
                residuals = y - X @ beta
                var = np.var(residuals)

                # BIC评分
                k = len(parents) + 1  # 参数数量（包括截距）
                score = -n / 2 * np.log(2 * np.pi * var) - n / 2 - k / 2 * np.log(n)
            except:
                score = -np.inf

        return score

    def _compute_score(self, data, graph):
        """计算整个图的评分"""
        n_vars = graph.shape[0]
        total_score = 0

        for node in range(n_vars):
            parents = set(np.where(graph[:, node] == 1)[0])
            total_score += self._compute_local_score(data, node, parents)

        return total_score

    def _is_dag(self, graph):
        """检查是否为有向无环图"""
        n = graph.shape[0]
        visited = [False] * n
        rec_stack = [False] * n

        def dfs(v):
            visited[v] = True
            rec_stack[v] = True

            for u in range(n):
                if graph[v, u] == 1:  # v -> u
                    if not visited[u]:
                        if dfs(u):
                            return True
                    elif rec_stack[u]:
                        return True

            rec_stack[v] = False
            return False

        for i in range(n):
            if not visited[i]:
                if dfs(i):
                    return False
        return True

    def fit(self, data):
        """运行GES算法"""
        n_vars = data.shape[1]

        # 初始化空图
        self.graph = np.zeros((n_vars, n_vars), dtype=int)
        current_score = self._compute_score(data, self.graph)

        # 前向阶段：添加边
        improved = True
        while improved:
            improved = False
            best_delta = 0
            best_edge = None

            for i in range(n_vars):
                for j in range(n_vars):
                    if i == j or self.graph[i, j] == 1:
                        continue

                    # 尝试添加边 i -> j
                    self.graph[i, j] = 1
                    if self._is_dag(self.graph):
                        new_score = self._compute_score(data, self.graph)
                        delta = new_score - current_score
                        if delta > best_delta:
                            best_delta = delta
                            best_edge = (i, j)
                    self.graph[i, j] = 0

            if best_edge is not None:
                self.graph[best_edge[0], best_edge[1]] = 1
                current_score += best_delta
                improved = True

        # 后向阶段：删除边
        improved = True
        while improved:
            improved = False
            best_delta = 0
            best_edge = None

            for i in range(n_vars):
                for j in range(n_vars):
                    if self.graph[i, j] == 0:
                        continue

                    # 尝试删除边 i -> j
                    self.graph[i, j] = 0
                    new_score = self._compute_score(data, self.graph)
                    delta = new_score - current_score
                    if delta > best_delta:
                        best_delta = delta
                        best_edge = (i, j)
                    self.graph[i, j] = 1

            if best_edge is not None:
                self.graph[best_edge[0], best_edge[1]] = 0
                current_score += best_delta
                improved = True

        return self

    def get_graph(self):
        return self.graph

三、因果推理模型实现

3.1 基于神经网络的结构因果模型

使用 MindSpore 实现可学习的结构因果模型，支持从数据中学习因果机制。

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Tensor, Parameter
from mindspore.common.initializer import Normal

class CausalMechanism(nn.Cell):
    """因果机制：使用神经网络建模函数关系"""

    def __init__(self, input_dim, hidden_dim=64):
        super().__init__()
        self.net = nn.SequentialCell([
            nn.Dense(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dense(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dense(hidden_dim, 1)
        ])

    def construct(self, x):
        return self.net(x)


class StructuralCausalModel(nn.Cell):
    """结构因果模型：可学习的因果图"""

    def __init__(self, n_vars, hidden_dim=64):
        super().__init__()
        self.n_vars = n_vars

        # 每个变量一个因果机制
        self.mechanisms = nn.CellList([
            CausalMechanism(n_vars, hidden_dim) 
            for _ in range(n_vars)
        ])

        # 可学习的因果图参数（注意力权重）
        self.adj_logits = Parameter(
            ms.numpy.zeros((n_vars, n_vars), dtype=ms.float32),
            requires_grad=True
        )

        # 噪声标准差
        self.noise_std = Parameter(
            ms.numpy.ones(n_vars, dtype=ms.float32) * 0.1,
            requires_grad=True
        )

    def get_adjacency(self, threshold=0.5):
        """获取因果图的邻接矩阵"""
        probs = ops.sigmoid(self.adj_logits)
        return (probs > threshold).astype(ms.float32)

    def construct(self, u_noise):
        """
        前向传播：根据噪声生成观测变量

        Args:
            u_noise: 外生噪声，shape=(batch_size, n_vars)
        """
        batch_size = u_noise.shape[0]

        # 使用消息传递按拓扑顺序计算
        # 这里简化处理，假设变量按因果顺序排列
        x = ops.zeros((batch_size, self.n_vars), dtype=ms.float32)

        # 获取邻接矩阵权重
        adj_weights = ops.sigmoid(self.adj_logits)

        for i in range(self.n_vars):
            # 计算父节点的影响
            parent_input = x * adj_weights[:, i].unsqueeze(0)
            parent_input = ops.cat([parent_input, u_noise[:, i:i+1]], axis=1)

            # 通过因果机制
            x[:, i] = self.mechanisms[i](parent_input).squeeze(1)

        return x

    def intervention(self, x_obs, intervene_var, intervene_value):
        """
        执行干预操作 do(X_i = value)

        Args:
            x_obs: 观测数据
            intervene_var: 被干预的变量索引
            intervene_value: 干预值
        """
        batch_size = x_obs.shape[0]
        x_intervened = x_obs.copy()

        # 设置干预值
        if isinstance(intervene_value, Tensor):
            x_intervened[:, intervene_var] = intervene_value
        else:
            x_intervened[:, intervene_var] = intervene_value

        # 重新计算下游变量
        adj_weights = ops.sigmoid(self.adj_logits)

        for j in range(intervene_var + 1, self.n_vars):
            parent_input = x_intervened * adj_weights[:, j].unsqueeze(0)
            noise = ops.randn((batch_size, 1), dtype=ms.float32) * ops.relu(self.noise_std[j])
            parent_input = ops.cat([parent_input, noise], axis=1)
            x_intervened[:, j] = self.mechanisms[j](parent_input).squeeze(1)

        return x_intervened


class CausalModelTrainer:
    """因果模型训练器"""

    def __init__(self, scm, learning_rate=0.001):
        self.scm = scm
        self.optimizer = nn.Adam(scm.trainable_params(), learning_rate=learning_rate)

    def compute_loss(self, x_obs, u_noise):
        """计算重建损失"""
        x_pred = self.scm(u_noise)
        loss = ops.mean((x_pred - x_obs) ** 2)

        # 稀疏性正则化（鼓励稀疏因果图）
        sparsity_reg = 0.01 * ops.mean(ops.sigmoid(self.scm.adj_logits))

        return loss + sparsity_reg

    def train_step(self, x_obs, u_noise):
        """单步训练"""
        grad_fn = ms.ops.value_and_grad(self.compute_loss, None, self.scm.trainable_params())
        loss, grads = grad_fn(x_obs, u_noise)
        self.optimizer(grads)
        return loss

    def fit(self, data, n_epochs=100, batch_size=32):
        """
        训练因果模型

        Args:
            data: 观测数据，shape=(n_samples, n_vars)
        """
        n_samples = data.shape[0]
        n_vars = data.shape[1]

        dataset = ms.dataset.NumpySlicesDataset(data, shuffle=True)
        dataset = dataset.batch(batch_size)

        for epoch in range(n_epochs):
            total_loss = 0
            n_batches = 0

            for batch in dataset:
                x_obs = Tensor(batch[0], dtype=ms.float32)
                u_noise = ops.randn((x_obs.shape[0], n_vars), dtype=ms.float32)

                loss = self.train_step(x_obs, u_noise)
                total_loss += loss.asnumpy()
                n_batches += 1

            if (epoch + 1) % 10 == 0:
                print(f'Epoch {epoch+1}/{n_epochs}, Loss: {total_loss/n_batches:.4f}')

        return self.scm


# 示例：训练因果模型
def demo_causal_model():
    """演示结构因果模型的训练"""
    np.random.seed(42)

    # 生成模拟因果数据
    n_samples = 1000
    U1 = np.random.randn(n_samples)
    U2 = np.random.randn(n_samples)
    U3 = np.random.randn(n_samples)

    X1 = U1
    X2 = 0.5 * X1 + U2
    X3 = 0.3 * X1 + 0.7 * X2 + U3

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

    # 创建并训练因果模型
    scm = StructuralCausalModel(n_vars=3, hidden_dim=32)
    trainer = CausalModelTrainer(scm, learning_rate=0.01)

    trained_scm = trainer.fit(data, n_epochs=50, batch_size=64)

    # 查看学习到的因果图
    adj = trained_scm.get_adjacency(threshold=0.3)
    print("\n学习到的因果图邻接矩阵：")
    print(adj.asnumpy())

    return trained_scm


if __name__ == '__main__':
    demo_causal_model()

3.2 反事实推理实现

反事实推理回答"如果当时 X 取不同的值，Y 会是多少"这类问题。

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Tensor

class CounterfactualReasoner:
    """反事实推理器"""

    def __init__(self, scm):
        """
        Args:
            scm: 训练好的结构因果模型
        """
        self.scm = scm

    def abduction(self, x_obs):
        """
        溯因：根据观测推断外生噪声

        给定观测 X=x，推断 U 使得 SCM(U) = x
        """
        # 简化实现：假设加性噪声模型
        # U_i = X_i - f_i(PA_i)
        batch_size = x_obs.shape[0]
        n_vars = x_obs.shape[1]

        u_inferred = ops.zeros((batch_size, n_vars), dtype=ms.float32)
        adj_weights = ops.sigmoid(self.scm.adj_logits)

        for i in range(n_vars):
            parent_input = x_obs * adj_weights[:, i].unsqueeze(0)
            parent_input = ops.cat([parent_input, ops.zeros((batch_size, 1))], axis=1)
            predicted = self.scm.mechanisms[i](parent_input).squeeze(1)
            u_inferred[:, i] = x_obs[:, i] - predicted

        return u_inferred

    def intervention_step(self, x, intervene_var, intervene_value):
        """执行单步干预"""
        x_new = x.copy()
        if isinstance(intervene_value, Tensor):
            x_new[:, intervene_var] = intervene_value
        else:
            x_new[:, intervene_var] = intervene_value
        return x_new

    def counterfactual(self, x_obs, intervene_var, intervene_value):
        """
        反事实推理

        Args:
            x_obs: 实际观测的事实
            intervene_var: 假设干预的变量
            intervene_value: 假设的干预值

        Returns:
            反事实结果：在干预下本会观测到的值
        """
        # Step 1: 溯因 - 推断外生噪声
        u_noise = self.abduction(x_obs)

        # Step 2: 干预 - 修改因果机制
        # Step 3: 预测 - 使用修改后的机制和推断的噪声预测

        batch_size = x_obs.shape[0]
        n_vars = x_obs.shape[1]

        # 初始化反事实结果
        x_cf = ops.zeros((batch_size, n_vars), dtype=ms.float32)
        adj_weights = ops.sigmoid(self.scm.adj_logits)

        for i in range(n_vars):
            if i == intervene_var:
                # 被干预变量直接取干预值
                if isinstance(intervene_value, Tensor):
                    x_cf[:, i] = intervene_value
                else:
                    x_cf[:, i] = intervene_value
            else:
                # 其他变量按因果机制计算
                parent_input = x_cf * adj_weights[:, i].unsqueeze(0)
                noise = u_noise[:, i:i+1] * ops.relu(self.scm.noise_std[i])
                parent_input = ops.cat([parent_input, noise], axis=1)
                x_cf[:, i] = self.scm.mechanisms[i](parent_input).squeeze(1)

        return x_cf

    def average_treatment_effect(self, data, treatment_var, outcome_var, 
                                  treatment_value, control_value):
        """
        计算平均处理效应（ATE）

        ATE = E[Y | do(X=1)] - E[Y | do(X=0)]
        """
        n_samples = data.shape[0]

        # 干预处理
        x_treated = data.copy()
        x_treated[:, treatment_var] = treatment_value

        x_control = data.copy()
        x_control[:, treatment_var] = control_value

        # 计算反事实结果
        cf_treated = self.counterfactual(data, treatment_var, treatment_value)
        cf_control = self.counterfactual(data, treatment_var, control_value)

        # ATE
        ate = ops.mean(cf_treated[:, outcome_var] - cf_control[:, outcome_var])

        return ate.asnumpy()


def demo_counterfactual():
    """演示反事实推理"""
    np.random.seed(42)

    # 简单因果模型：X -> Y
    # Y = 2*X + noise
    n_samples = 100
    X = np.random.randn(n_samples)
    noise = np.random.randn(n_samples) * 0.5
    Y = 2 * X + noise

    data = np.column_stack([X, Y])
    data_tensor = Tensor(data, dtype=ms.float32)

    # 创建并训练SCM
    scm = StructuralCausalModel(n_vars=2, hidden_dim=16)
    trainer = CausalModelTrainer(scm, learning_rate=0.01)
    trained_scm = trainer.fit(data, n_epochs=30, batch_size=32)

    # 反事实推理
    reasoner = CounterfactualReasoner(trained_scm)

    # 计算ATE：如果X从0变为1，Y会如何变化
    ate = reasoner.average_treatment_effect(
        data_tensor, 
        treatment_var=0, 
        outcome_var=1,
        treatment_value=1.0,
        control_value=0.0
    )

    print(f"\n平均处理效应 (ATE): {ate:.4f}")
    print(f"真实因果效应: 2.0 (因为 Y = 2*X + noise)")

    return reasoner


if __name__ == '__main__':
    demo_counterfactual()

四、实际应用案例：医疗诊断中的因果推理

4.1 问题背景

在医疗诊断中，我们需要区分症状与疾病的因果关系，避免混淆相关性与因果性。例如：

• 症状（发烧、咳嗽）与疾病（流感、肺炎）的关系
• 治疗方案与康复结果的因果效应
• 潜在混杂因素（年龄、基础疾病）的影响

4.2 完整实现

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Tensor
import numpy as np

class MedicalCausalModel:
    """医疗诊断因果推理模型"""

    def __init__(self):
        # 变量定义
        self.var_names = [
            'age',        # 年龄（标准化）
            'smoking',    # 吸烟史（0/1）
            'exercise',   # 运动频率（标准化）
            'blood_pressure',  # 血压（标准化）
            'cholesterol',     # 胆固醇（标准化）
            'heart_disease',   # 心脏病（0/1）
            'treatment',       # 治疗方案（0/1）
            'recovery'         # 康复情况（标准化）
        ]

        # 预定义因果结构
        # 年龄 -> 血压, 胆固醇, 心脏病
        # 吸烟 -> 血压, 心脏病
        # 运动 -> 血压, 胆固醇
        # 血压 -> 心脏病
        # 胆固醇 -> 心脏病
        # 心脏病 -> 治疗
        # 治疗 -> 康复
        # 心脏病 -> 康复

        self.causal_graph = self._build_causal_graph()
        self.mechanisms = self._build_mechanisms()

    def _build_causal_graph(self):
        """构建因果图邻接矩阵"""
        n = len(self.var_names)
        graph = np.zeros((n, n), dtype=int)

        # 定义边
        edges = [
            (0, 3), (0, 4), (0, 5),  # age -> bp, chol, hd
            (1, 3), (1, 5),           # smoking -> bp, hd
            (2, 3), (2, 4),           # exercise -> bp, chol
            (3, 5),                   # bp -> hd
            (4, 5),                   # chol -> hd
            (5, 6),                   # hd -> treatment
            (6, 7),                   # treatment -> recovery
            (5, 7),                   # hd -> recovery
        ]

        for i, j in edges:
            graph[i, j] = 1

        return graph

    def _build_mechanisms(self):
        """构建因果机制"""
        # 使用线性模型简化
        mechanisms = {
            'blood_pressure': {
                'parents': ['age', 'smoking', 'exercise'],
                'coefficients': {'age': 0.3, 'smoking': 0.4, 'exercise': -0.2},
                'intercept': 0.0
            },
            'cholesterol': {
                'parents': ['age', 'exercise'],
                'coefficients': {'age': 0.25, 'exercise': -0.15},
                'intercept': 0.0
            },
            'heart_disease': {
                'parents': ['age', 'smoking', 'blood_pressure', 'cholesterol'],
                'coefficients': {'age': 0.2, 'smoking': 0.3, 'blood_pressure': 0.35, 'cholesterol': 0.25},
                'intercept': -0.5
            },
            'treatment': {
                'parents': ['heart_disease'],
                'coefficients': {'heart_disease': 0.8},
                'intercept': 0.0
            },
            'recovery': {
                'parents': ['heart_disease', 'treatment'],
                'coefficients': {'heart_disease': -0.4, 'treatment': 0.6},
                'intercept': 0.5
            }
        }
        return mechanisms

    def simulate_data(self, n_samples=1000, seed=42):
        """模拟医疗数据"""
        np.random.seed(seed)

        # 外生变量
        age = np.random.randn(n_samples)  # 标准化年龄
        smoking = np.random.binomial(1, 0.3, n_samples)  # 30%吸烟
        exercise = np.random.randn(n_samples)  # 标准化运动频率

        # 内生变量（按因果顺序）
        noise_bp = np.random.randn(n_samples) * 0.3
        blood_pressure = (0.3 * age + 0.4 * smoking - 0.2 * exercise + noise_bp)

        noise_chol = np.random.randn(n_samples) * 0.3
        cholesterol = (0.25 * age - 0.15 * exercise + noise_chol)

        noise_hd = np.random.randn(n_samples) * 0.3
        hd_score = (0.2 * age + 0.3 * smoking + 0.35 * blood_pressure + 
                    0.25 * cholesterol - 0.5 + noise_hd)
        heart_disease = (hd_score > 0).astype(float)

        # 治疗决策（有心脏病则更可能接受治疗）
        treatment_prob = 0.2 + 0.6 * heart_disease
        treatment = np.random.binomial(1, treatment_prob)

        # 康复情况
        noise_rec = np.random.randn(n_samples) * 0.3
        recovery = (-0.4 * heart_disease + 0.6 * treatment + 0.5 + noise_rec)

        data = np.column_stack([
            age, smoking, exercise, blood_pressure, cholesterol,
            heart_disease, treatment, recovery
        ])

        return data

    def compute_causal_effect(self, data, treatment_var, outcome_var):
        """
        计算因果效应（使用后门调整）

        P(Y | do(T)) = Σ_z P(Y | T, Z) * P(Z)
        """
        # 识别混杂因素
        # 对于 treatment -> recovery，混杂因素是 heart_disease
        confounder_idx = 5  # heart_disease

        # 分层计算
        treatment_vals = [0, 1]
        outcome_means = []

        for t_val in treatment_vals:
            # 按混杂因素分层
            weighted_sum = 0
            for z_val in [0, 1]:
                mask = (data[:, confounder_idx] == z_val) & (data[:, treatment_var] == t_val)
                if mask.sum() > 0:
                    y_given_t_z = data[mask, outcome_var].mean()
                    p_z = (data[:, confounder_idx] == z_val).mean()
                    weighted_sum += y_given_t_z * p_z
            outcome_means.append(weighted_sum)

        # 因果效应
        causal_effect = outcome_means[1] - outcome_means[0]
        return causal_effect

    def backdoor_adjustment(self, data, treatment_idx, outcome_idx, adjustment_set):
        """
        后门调整公式实现

        P(Y | do(X)) = Σ_z P(Y | X, Z) * P(Z)
        """
        treatment_vals = np.unique(data[:, treatment_idx])

        results = {}
        for t_val in treatment_vals:
            adjusted_mean = 0

            # 遍历调整集的所有组合
            adjustment_vals = [np.unique(data[:, z]) for z in adjustment_set]

            from itertools import product
            for z_combo in product(*adjustment_vals):
                # 构建条件掩码
                mask = (data[:, treatment_idx] == t_val)
                for i, z_idx in enumerate(adjustment_set):
                    mask &= (data[:, z_idx] == z_combo[i])

                if mask.sum() > 0:
                    # P(Y | X=x, Z=z)
                    y_mean = data[mask, outcome_idx].mean()

                    # P(Z=z)
                    z_mask = np.ones(len(data), dtype=bool)
                    for i, z_idx in enumerate(adjustment_set):
                        z_mask &= (data[:, z_idx] == z_combo[i])
                    p_z = z_mask.mean()

                    adjusted_mean += y_mean * p_z

            results[t_val] = adjusted_mean

        return results


def demo_medical_causal():
    """演示医疗诊断因果推理"""
    print("=" * 60)
    print("医疗诊断因果推理示例")
    print("=" * 60)

    # 创建模型
    model = MedicalCausalModel()

    # 生成数据
    data = model.simulate_data(n_samples=5000)

    print("\n变量说明：")
    for i, name in enumerate(model.var_names):
        print(f"  [{i}] {name}")

    print("\n因果图边：")
    edges = np.where(model.causal_graph == 1)
    for i, j in zip(edges[0], edges[1]):
        print(f"  {model.var_names[i]} -> {model.var_names[j]}")

    # 计算治疗的因果效应
    print("\n" + "=" * 60)
    print("治疗对康复的因果效应分析")
    print("=" * 60)

    # 使用后门调整
    treatment_idx = 6  # treatment
    outcome_idx = 7    # recovery
    adjustment_set = [5]  # heart_disease 作为混杂因素

    results = model.backdoor_adjustment(data, treatment_idx, outcome_idx, adjustment_set)

    print(f"\n后门调整结果：")
    print(f"  E[Recovery | do(Treatment=0)] = {results[0]:.4f}")
    print(f"  E[Recovery | do(Treatment=1)] = {results[1]:.4f}")
    print(f"  因果效应 = {results[1] - results[0]:.4f}")

    # 对比朴素估计（有偏差）
    naive_effect = (data[data[:, 6] == 1, 7].mean() - 
                    data[data[:, 6] == 0, 7].mean())
    print(f"\n朴素估计（未调整混杂因素）= {naive_effect:.4f}")
    print(f"偏差 = {naive_effect - (results[1] - results[0]):.4f}")

    return model, data


if __name__ == '__main__':
    demo_medical_causal()

五、因果推理与机器学习的结合

5.1 因果正则化

将因果知识融入机器学习模型，提高泛化能力和可解释性。

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Tensor

class CausalRegularizedModel(nn.Cell):
    """因果正则化模型"""

    def __init__(self, input_dim, hidden_dim=64, causal_graph=None):
        super().__init__()

        self.encoder = nn.SequentialCell([
            nn.Dense(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dense(hidden_dim, hidden_dim),
            nn.ReLU()
        ])

        self.decoder = nn.Dense(hidden_dim, 1)

        # 因果图约束
        if causal_graph is not None:
            self.register_buffer('causal_mask', Tensor(causal_graph, dtype=ms.float32))
        else:
            self.causal_mask = None

        # 因果权重
        self.causal_weights = Parameter(
            ms.numpy.randn(input_dim, hidden_dim) * 0.1,
            requires_grad=True
        )

    def construct(self, x):
        # 编码
        h = self.encoder(x)

        # 因果约束的解码
        if self.causal_mask is not None:
            # 只使用因果相关的特征
            weighted_input = ops.matmul(x, self.causal_weights)
            h = h + weighted_input

        # 解码
        output = self.decoder(h)
        return output

    def causal_loss(self, x, y):
        """计算因果正则化损失"""
        pred = self.construct(x)

        # 预测损失
        pred_loss = ops.mean((pred - y) ** 2)

        # 因果稀疏性损失
        if self.causal_mask is not None:
            sparsity_loss = ops.mean(ops.abs(self.causal_weights * (1 - self.causal_mask)))
        else:
            sparsity_loss = ops.mean(ops.abs(self.causal_weights))

        return pred_loss + 0.1 * sparsity_loss


class InvariantRiskMinimization(nn.Cell):
    """不变风险最小化（IRM）"""

    def __init__(self, input_dim, hidden_dim=64):
        super().__init__()

        self.feature_extractor = nn.SequentialCell([
            nn.Dense(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dense(hidden_dim, hidden_dim),
            nn.ReLU()
        ])

        self.classifier = nn.Dense(hidden_dim, 1)

        # IRM惩罚的虚拟标量
        self.irm_lambda = 1.0

    def construct(self, x):
        features = self.feature_extractor(x)
        output = self.classifier(features)
        return output

    def irm_penalty(self, x, y):
        """计算IRM惩罚项"""
        features = self.feature_extractor(x)

        # 计算梯度惩罚
        # 简化实现：使用特征与残差的相关性
        pred = self.classifier(features)
        residual = pred - y

        # 惩罚特征与残差的相关性
        penalty = ops.mean(features * residual) ** 2

        return penalty

    def total_loss(self, x, y):
        """总损失 = 预测损失 + IRM惩罚"""
        pred = self.construct(x)
        pred_loss = ops.mean((pred - y) ** 2)
        irm_penalty = self.irm_penalty(x, y)

        return pred_loss + self.irm_lambda * irm_penalty


def demo_causal_ml():
    """演示因果机器学习"""
    np.random.seed(42)

    # 生成数据：两个环境
    n = 500

    # 环境1
    X1_env1 = np.random.randn(n)
    X2_env1 = 2 * X1_env1 + np.random.randn(n) * 0.5
    Y_env1 = X1_env1 + np.random.randn(n) * 0.3

    # 环境2（X2的分布改变）
    X1_env2 = np.random.randn(n)
    X2_env2 = 2 * X1_env2 + np.random.randn(n) * 2.0  # 更大噪声
    Y_env2 = X1_env2 + np.random.randn(n) * 0.3

    data_env1 = np.column_stack([X1_env1, X2_env1, Y_env1])
    data_env2 = np.column_stack([X1_env2, X2_env2, Y_env2])

    print("因果机器学习示例：不变风险最小化")
    print("=" * 50)
    print("\n真实因果结构：X1 -> Y, X1 -> X2")
    print("X2 是虚假相关特征（不应被用于预测）")

    # 创建IRM模型
    irm = InvariantRiskMinimization(input_dim=2, hidden_dim=32)

    # 训练（简化）
    print("\n训练IRM模型...")
    print("模型将学习只使用因果特征 X1 进行预测")

    return irm


if __name__ == '__main__':
    demo_causal_ml()

六、总结与展望

6.1 本文要点

1. 因果发现：PC算法和GES算法可以从观测数据中学习因果结构
2. 因果推理：通过do-演算和反事实推理回答干预问题
3. 结构因果模型：使用神经网络建模可学习的因果机制
4. 实际应用：医疗诊断中的因果效应估计和混杂因素调整
5. 因果机器学习：因果正则化和不变风险最小化提高模型泛化能力

6.2 MindSpore 优势

• 自动微分：简化因果模型中复杂梯度的计算
• 分布式训练：支持大规模因果发现算法
• 灵活架构：便于实现各种因果推理模型

6.3 扩展阅读

• Judea Pearl《因果论》
• Jonas Peters《因果推断基础》
• MindSpore 官方文档

因果推理是通往强人工智能的关键技术，掌握它将帮助我们构建更可靠、更可解释的AI系统。