Maybe there's a bug in 'apply_unitary_density_bmm' function?

[TuDouuuuu]

https://github.com/mit-han-lab/torchquantum/blob/c8a1ab64e2c0e9a6b570e2307e5ad10a5560f2b1/torchquantum/functional/gate_wrapper.py
LIne 301

    """
    Compute \rho U^\dagger 
    """
    matdag = torch.conj(mat)
    matdag = matdag.type(C_DTYPE).to(density.device)

maybe matdag = torch.conj(mat) is not $U^\dagger$.

[TuDouuuuu]

I use $\ket{\psi}=\ket{0} \oplus \ket{0}$, and apply Y gate in $q0$.

import torch
import numpy as np
import torchquantum as tq
import torchquantum.functional as tqf


def apply_unitary_density_bmm(density, mat, wires):
    """Apply the unitary to the DensityMatrix using torch.bmm method.
    Args:
        state (torch.Tensor): The statevector.
        mat (torch.Tensor): The unitary matrix of the operation.
        wires (int or List[int]): Which qubit the operation is applied to.
    Returns:
        torch.Tensor: The new statevector.
    """
    device_wires = wires
    n_qubit = density.dim() // 2
    mat = mat.type(torch.complex64).to(density.device)
    """
    Compute U \rho
    """
    devices_dims = [w + 1 for w in device_wires]
    permute_to = list(range(density.dim()))
    for d in sorted(devices_dims, reverse=True):
        del permute_to[d]
    permute_to = permute_to[:1] + devices_dims + permute_to[1:]
    permute_back = list(np.argsort(permute_to))
    original_shape = density.shape
    permuted = density.permute(permute_to).reshape(
        [original_shape[0], mat.shape[-1], -1])

    if len(mat.shape) > 2:
        # both matrix and state are in batch mode
        new_density = mat.bmm(permuted)
    else:
        # matrix no batch, state in batch mode
        bsz = permuted.shape[0]
        expand_shape = [bsz] + list(mat.shape)
        new_density = mat.expand(expand_shape).bmm(permuted)
    new_density = new_density.view(original_shape).permute(permute_back)
    """
    Compute \rho U^\dagger 
    """
    matdag = torch.conj(mat)
    matdag = matdag.type(torch.complex64).to(density.device)

    devices_dims_dag = [n_qubit + w + 1 for w in device_wires]
    permute_to_dag = list(range(density.dim()))
    for d in sorted(devices_dims_dag, reverse=True):
        del permute_to_dag[d]
    permute_to_dag = permute_to_dag + devices_dims_dag
    permute_back_dag = list(np.argsort(permute_to_dag))
    permuted_dag = new_density.permute(permute_to_dag).reshape(
        [original_shape[0], -1, matdag.shape[0]])

    if len(matdag.shape) > 2:
        # both matrix and state are in batch mode
        new_density = permuted_dag.bmm(matdag)
    else:
        # matrix no batch, state in batch mode
        bsz = permuted_dag.shape[0]
        expand_shape = [bsz] + list(matdag.shape)
        new_density = permuted_dag.bmm(matdag.expand(expand_shape))
    new_density = new_density.view(original_shape).permute(permute_back_dag)
    return new_density


density = torch.zeros(2, 2, 2, 2,  dtype=torch.complex64)
density[0,0,0,0] = 1
mat = tq.PauliY().matrix
print(apply_unitary_density_bmm(density=density, mat=mat, wires=[0]))

The output is :

tensor([[[[ 0.+0.j,  0.+0.j],
          [ 0.+0.j,  0.+0.j]],

         [[ 0.+0.j, -1.+0.j],
          [ 0.+0.j,  0.+0.j]]],


        [[[ 0.+0.j,  0.+0.j],
          [ 0.+0.j,  0.+0.j]],

         [[ 0.+0.j,  0.+0.j],
          [ 0.+0.j,  0.+0.j]]]])

But after applying Y gate, maybe the density is $diag(0, 1, 0, 0)$.

[TuDouuuuu]

In https://github.com/mit-han-lab/torchquantum/blob/f1c3dfe01d7f347e80211bbac5f1369bf49876d4/torchquantum/functional/gate_wrapper.py , it seems right.
But in https://github.com/mit-han-lab/torchquantum/blob/f1c3dfe01d7f347e80211bbac5f1369bf49876d4/torchquantum/density/density_func.py, it maybe still has something wrong.

[GenericP3rson]

Hi, thanks for letting us know! We are actively working on rebooting our density matrix in the branch dev-zhuoyang-dm and I've added @yezhuoyang to the issue to keep him in the loop.

Feel free to open a PR to that branch as well so it's integrated in the next DM release!