在Python中计算外部产品的部分痕迹？

计算科学 线性代数 Python 量子力学
2021-12-01 06:49:35

我有一个计算选定维度上的部分跟踪的 python 实现。

def partial_trace(rho, keep, dims):
    """Calculate the partial trace

    Parameters
    ----------
    ρ : 2D array
        Matrix to trace
    keep : array
        An array of indices of the spaces to keep after
        being traced. For instance, if the space is
        A x B x C x D and we want to trace out B and D,
        keep = [0,2]
    dims : array
        An array of the dimensions of each space.
        For instance, if the space is A x B x C x D,
        dims = [dim_A, dim_B, dim_C, dim_D]

    Returns
    -------
    ρ_a : 2D array
        Traced matrix
    """
    dims = np.asarray(dims)
    N = dims.size
    # Indices to trace
    rest = np.delete(np.arange(N), keep)
    # Reshape into tensor
    rho_a = rho.reshape(np.tile(dims, 2))
    # Trace indices
    for i,d in enumerate(rest[::-1]):
        rho_a = np.trace(rho_a, axis1=d, axis2=N-i+d)
    # Reshape into matrix
    N = np.prod(dims[keep])
    return rho_a.reshape(N,N)

我想计算部分痕迹|u⟩⟨u|但是当我尝试构建外部产品时遇到了内存不足错误。相反，我想直接从|u⟩.

我将如何在 Python 中做到这一点？

1个回答

我修改了部分跟踪的实现以使用 einsum。

def partial_trace(rho, keep, dims, optimize=False):
    """Calculate the partial trace

    ρ_a = Tr_b(ρ)

    Parameters
    ----------
    ρ : 2D array
        Matrix to trace
    keep : array
        An array of indices of the spaces to keep after
        being traced. For instance, if the space is
        A x B x C x D and we want to trace out B and D,
        keep = [0,2]
    dims : array
        An array of the dimensions of each space.
        For instance, if the space is A x B x C x D,
        dims = [dim_A, dim_B, dim_C, dim_D]

    Returns
    -------
    ρ_a : 2D array
        Traced matrix
    """
    keep = np.asarray(keep)
    dims = np.asarray(dims)
    Ndim = dims.size
    Nkeep = np.prod(dims[keep])

    idx1 = [i for i in range(Ndim)]
    idx2 = [Ndim+i if i in keep else i for i in range(Ndim)]
    rho_a = rho.reshape(np.tile(dims,2))
    rho_a = np.einsum(rho_a, idx1+idx2, optimize=optimize)
    return rho_a.reshape(Nkeep, Nkeep)

这可以修改为将向量作为输入。

def ptrace_outer(u, keep, dims, optimize=False):
    """Calculate the partial trace of an outer product

    ρ_a = Tr_b(|u><u|)

    Parameters
    ----------
    u : array
        Vector to use for outer product
    keep : array
        An array of indices of the spaces to keep after
        being traced. For instance, if the space is
        A x B x C x D and we want to trace out B and D,
        keep = [0,2]
    dims : array
        An array of the dimensions of each space.
        For instance, if the space is A x B x C x D,
        dims = [dim_A, dim_B, dim_C, dim_D]

    Returns
    -------
    ρ_a : 2D array
        Traced matrix
    """
    keep = np.asarray(keep)
    dims = np.asarray(dims)
    Ndim = dims.size
    Nkeep = np.prod(dims[keep])

    idx1 = [i for i in range(Ndim)]
    idx2 = [Ndim+i if i in keep else i for i in range(Ndim)]
    u = u.reshape(dims)
    rho_a = np.einsum(u, idx1, u.conj(), idx2, optimize=optimize)
    return rho_a.reshape(Nkeep, Nkeep)
其它你可能感兴趣的问题
上一篇防止 FEM 模拟中的单节点尖峰（使用连续 Galerkin） 下一篇求解具有 Neumann 边界的 3-D 热方程