矩阵求偏迹（约化密度矩阵）python代码（2）

这篇文章接着上一篇分享的求偏迹代码，继续更新一些实用的函数，希望对大家有帮助！本文依然是对 2n×2n 的矩阵进行求解。算法思想依旧是将大问题分解为小问题，递归。我们在pytorch框架下（大家也可以用AI转换到其他数据格式下），

首先定义矩阵分块和组合的操作：

def slice_matrix(matrix):
    """
    将输入的矩阵进行分块切片，返回四个子矩阵的列表
    """
    n = matrix.shape[0] // 2
    sub_matrices = [matrix[:n, :n], matrix[:n, n:], matrix[n:, :n], matrix[n:, n:]]
    return sub_matrices


def combine_matrix(sub_matrices):
    """
    输入存储四个矩阵的列表，将四个子矩阵按照特定方式拼接成一个大矩阵
    """
    top = torch.cat(sub_matrices[:2], dim=1)
    bottom = torch.cat(sub_matrices[2:], dim=1)
    return torch.cat((top, bottom), dim=0)

然后就是我们上一篇文章的矩阵偏迹代码：

def partial_trace(matrix, trace_index):
    """
    对矩阵进行部分迹操作，trace_index表示要求迹的空间，从0开始。
    """
    sub_matrices = slice_matrix(matrix)
    if trace_index == 0:
        return sub_matrices[0] + sub_matrices[3]
    else:
        traced_slices = [partial_trace(sub_matrix, trace_index - 1) for sub_matrix in s