I'm trying to derive equation $(1)$ on p.2 in Lloyd et al, 2013 which reads
$$ \text{Tr}_A\left[\exp(-i\theta S_{AB}) (\rho_A \otimes \sigma_B) \exp(i\theta S_{AB}) \right] = (\cos^2 \theta) \sigma_B + (\sin^2 \theta) \rho_B - i \sin \theta [\rho_B, \sigma_B] $$
where $S_{AB}$ is a swap operator for the two systems. My approach so far is to write the operator as $$ \exp(i\theta S_{AB}) = \begin{pmatrix} e^{i\theta} & 0 & 0 & 0 \\ 0 & \cos \theta & i \sin\theta & 0 \\ 0 & i \sin \theta & \cos \theta & 0 \\ 0 & 0 & 0 & e^{i\theta} \end{pmatrix} \begin{matrix} \leftarrow\rho_A \otimes \rho_B \\ \leftarrow\rho_A \otimes \sigma_B\\ \leftarrow\sigma_A \otimes \rho_B \\ \leftarrow\sigma_A\otimes \sigma_B \end{matrix} $$
where the labels describe how each input is affected by this operation if the possible state configurations are represented as a column vector, e.g. $\rho\otimes \sigma \sim \hat{e}_1$ would become $\cos\theta (\rho\otimes \sigma) + \sin \theta (\sigma \otimes \rho) \sim \cos\theta \hat{e}_1 + \sin \theta \hat{e}_2$.
However, I haven't been able to carry out the calculation since I'm unsure how to apply the operator or partial trace on $(\rho\otimes \sigma)$ when I set up operators like this. Is there anything pathological about this representation of a SWAP, and if so what is the right approach to deriving the above equation?