python - 用tensorflow计算一个量子算子的时间依赖期望值

Question

我正在尝试计算量子算子随时间变化的期望值。是的，我得到了每一次的状态

t_l = tf.constant( np.linspace(0, 4 / gam_q, 1000, dtype = np.complex128) )
dt = t_l[1] - t_l[0]

H_t = tf.tensordot(- 1j * t_l, H, 0)
H_t = tfla.expm(H_t)

psi0 = tf.constant( ((qtp.fock(2, 0)).unit()).full() )

psi_t = tf.tensordot(H_t, psi0, [[2], [0]])

其中 H 是哈密顿量（在这种情况下为 sigma_z）。但是，我需要知道算子 S_op 的期望值的时间依赖性。我试过这个

pre = tf.tensordot( tf.math.conj(psi_t), S_op, [[1], [0]]) #psi_t was not transposed, so I put [[1], [0]]
expect = tf.tensordot(pre, psi_t, [[2], [1]])

它返回一个形状为 (1000, 1, 1000, 1) 的张量，而我需要一个形状为 (1000) 的一维数组。

编辑

我通过引入 einsum 解决了

pre = tf.tensordot(tfmt.conj(psi_t), S_op, [[1], [0]])
expect = tfmt.real( tf.tensordot(pre, psi_t, [[2], [1]]) )
expect = tf.einsum("ijij->i", expect)

Answer 1

为了社区的利益，在答案部分提供解决方案

我通过引入解决了einsum

pre = tf.tensordot(tfmt.conj(psi_t), S_op, [[1], [0]])
expect = tfmt.real( tf.tensordot(pre, psi_t, [[2], [1]]) )
expect = tf.einsum("ijij->i", expect)

（从 Alberto Mercurio 转述）