正如注解中提到的，您复制了不必要的输入数组。此外，如您所建议的，可以直接使用np.einsum计算外积的部分迹线，而无需分配C。但这对代码的可读性没有帮助。..
看看ptrace的实现和你传递给它的参数（也许下次尝试将这样的外部函数提炼成它的Numpy等价物-通过一些上下文可能更容易理解），看起来你最终会调用_ptrace_dense（as found here），特别是该函数的else部分：

rhomat = np.trace(Q.full()
    .reshape(rd + rd)
    .transpose(qtrace + [nd + q for q in qtrace] +
                sel + [nd + q for q in sel])
    .reshape([np.prod(dtrace, dtype=np.int32),
            np.prod(dtrace, dtype=np.int32),
            np.prod(dkeep, dtype=np.int32),
            np.prod(dkeep, dtype=np.int32)]))

正如你所看到的，这做了一些复杂的整形/换位，然后取结果的np.trace。您可以利用外积的踪迹实际上是内积的事实，这两者正是np.einsum的目的。这消除了大量的乘法运算，并且需要少得多的存储器。
您还可以对np.einsum中的最终尺寸求和，然后将输出除以最终尺寸的长度，以消除.mean(axis=3)。
为了在一个np.einsum步骤中完成所有操作，您必须首先重塑一个输入数组。最后可以归结为：

A = A.reshape([4, 8, 4, 8, len(t_list), Mmax])

X = np.einsum('ikok...q,ll...q->io...', A, B).reshape((4, 4, len(t_list))) / Mmax

X现在保存for循环的堆栈结果，即：e.所有ptrace结果在最终轴上。
这种转变不是特别直观，至少对我来说是这样。然而，到达这里的步骤相当容易，如果冗长。
我将通过我在特定情况下所做的工作：你有两个形状为（32，32）的输入数组，想要前两个子集/量子位/任何东西的部分轨迹。在你的其他维度上传播这个操作应该是显而易见的（如果不是，请阅读here）。
总结一下你要做的事情，它可以归结为：

C = np.einsum('ij,kl->ikjl',A,B)
# While you reshape this output, this has no effect because internally `ptrace` reshapes it internally
X = C.reshape([2, 2, 2, ..., 2]) # Reshape it into a 2x2x...x2 array
Y = X.transpose([2, 3, ..., 9, 12, 13, ..., 19, 0, 1, 10, 11]) # Move the first two axes of the first and second original axes to the end
Z = Y.reshape([256, 256, 4, 4])
output = np.trace(Z)

使用axis1和axis2可以显著简化这一点（不知道为什么Qutip库不这样做）：

C = np.einsum('ij,kl->ikjl',A,B)
X = C.reshape([4, 256, 4, 256])
output = np.trace(X, axis1=1, axis2=3)

更换np.trace：

C = np.einsum('ij,kl->ikjl', A, B)
X = C.reshape([4, 256, 4, 256])
output = np.einsum('ijkj->ik', X)

两个稍微不直观的东西可能有助于解释接下来的几个步骤：

A = np.random.rand(4, 4)
B = np.random.rand(4, 4)
X = np.einsum('ij,kl->ijkl', A, B)  # Note this is ijkl NOT ikjl
Y = np.einsum('i,j->ij', A.flatten(), B.flatten())
assert np.allclose(X.flatten(), Y.flatten())

（虽然不是很重要）：

A = np.random.rand(4, 4)
B = np.random.rand(4, 4)
X = np.einsum('ij,kl->ikjl', A, B)
Y = np.einsum('ij,kl->ijkl', A, B).transpose((0, 2, 1, 3)) # Or .swapaxes(1, 2)
assert np.allclose(X.flatten(), Y.flatten())

第一点的意思是：

A = np.random.rand(32, 32)
B = np.random.rand(32, 32)
C = np.einsum('ij,kl->ikjl', A, B)
A = A.reshape((4, 8, 4, 8))
B = B.reshape((4, 8, 4, 8))
D = np.einsum('ij kl,mn op->ij mn kl op', A, B) # Spaces have no meaning, just help map the dimension groups to previous names,
# So i=ij, j=kl, k=mn and l=op
assert np.allclose(C.flatten(), D.flatten())

这很有帮助，因为为了合并两个np.einsum，你需要能够将A和B（好吧，也许不是B，稍后会详细介绍）拆分为32 x32的第一个np.einsum和4x 256的第二个，如前所示。因此，A和B的形状必须是（4，8，4，8），因为4*8=32满足第一个要求，而8*4*8=256满足第二个要求。这就把以前的问题简化为：

A = np.random.rand(4, 8, 4, 8)
B = np.random.rand(4, 8, 4, 8)
C = np.einsum('ij kl,mn op->ij mn kl op', A, B)
output = np.einsum('i jmn k jmn->ik', C).reshape((4, 4))

现在你要做的就是合并两个np.einsum。这很容易，因为第一个输出的ijmnklop直接Map到第二个输入的ijmnkjmn。因此：

A = np.random.rand(4, 8, 4, 8)
B = np.random.rand(4, 8, 4, 8)
output = np.einsum('i j kj,m n mn->ik', A, B).reshape((4, 4))

这里应该很明显，B实际上不需要4维，只需要2维。最后得出：

A = np.random.rand(4, 8, 4, 8)
B = np.random.rand(32, 32)
output = np.einsum('i j kj,l l->ik', A, B).reshape((4, 4))

显然，与我一样使用np.random，不如将数据重塑为正确的形状。
这不是一个特别令人满意的解释，因为输出代码似乎与输入代码完全无关。但是，我希望这些步骤可以帮助您理解如何处理它，并将其推广到其他np.einsum案例。
关于泛化到其他子集的具体问题，您可以将A重塑为2x 2x。..x2并从中选择任意索引。考虑到前面的（意想不到的）简化，您不必修改B。例如，如果您想要0和3而不是0和1：

A = A.reshape((2,)*10)
B = B.reshape((32, 32))
output = np.einsum('i jk l m o jk p m,z z->i l o p', A, B).reshape((4, 4))

请注意，要“保留”的所有尺寸标注都使用唯一的字母命名，并在输出中重复，而要跟踪的所有尺寸标注都在输入中重复，不会出现在输出中。B被简单地跟踪，而不管选择哪个维度。
最后，你可以写一个函数来自动生成下标：

def get_subscripts(keep_indices: List[int], length: int, start_at='b', trace_char='a') -> str:
    start = ord(start_at)
    in_subs = list(range(start, start + length))*2
    out_subs = [start + k for k in keep_indices]
    for i, k in enumerate(keep_indices):
        in_subs[length + k] = start + length + i
        out_subs.append(start + length + i)
    return ''.join(chr(x) for x in in_subs) + ',' + trace_char*2 + '->' + ''.join(chr(x) for x in out_subs)

subsets = [0, 1, 3]
A = A.reshape((2,)*10)
B = B.reshape((32, 32))
output = np.einsum(get_subscripts(subsets, 5), A, B).reshape((2**len(subsets),)*2)

有趣的问题！