gradient of partial trace of matrix function

Question

Let $U = \exp(-i(\epsilon_1A + \epsilon_2B))$, where $\epsilon_1, \epsilon_2 \in \mathbb{R}$ and $A, B$ are Hermitian matrices acting on Hilbert space $V \otimes W$ . There is a straightforward way to calculate $\partial \text{Tr}(UK)/\partial \epsilon_1$ for an arbitrary Hermitian matrix $K$ on Hilbert space $V \otimes W$. But I couldn't see a way to do the calculation when the trace is replaced to a partial trace, i.e. calculating $\partial \text{Tr}_W(UK)/\partial \epsilon_1$. If anyone could enlighten me how to approach this problem or points me to a relevant reference, I'd be very grateful. Thank you in advance.

Answer 1

$ \def\bbR#1{{\mathbb R}^{#1}} \def\bbC#1{{\mathbb C}^{#1}} \def\o{{\large\tt1}} \def\e{\varepsilon} \def\a{\alpha} \def\b{\beta} \def\d{\delta} \def\BB#1{\Big[#1\Big]} \def\LR#1{\left(#1\right)} \def\BR#1{\Big(#1\Big)} \def\op{\operatorname} \def\size#1{\op{size}\LR{#1}} \def\Diag#1{\op{Diag}\LR{#1}} \def\trace#1{\op{tr}\LR{#1}} \def\ptrace#1{\op{ptr}\LR{#1}} \def\p{\partial} \def\grad#1#2{\frac{\p #1}{\p #2}} \def\m#1{\left[\begin{array}{c}#1\end{array}\right]} \def\qiq{\quad\implies\quad} \def\Sj{\sum_{j=\o}^n} \def\Sk{\sum_{k=\o}^n} \def\c#1{\color{red}{#1}} $Assume the dimensions of the spaces and matrices to be $$\eqalign{ n,n &= \size{V} \qquad m,m = \size{W} \qquad p=mn \\ p,p &= \size{A} = \size{B} = \size{U} = \size{K} \\ }$$ Let $\e_i\in\bbR{n}$ denote the standard orthogonal basis vectors, and use the identity matrix $I\in\bbR{m\times m}$ to construct their block-matrix analogs $$ E_k = \e_k\otimes I \qiq E_j^TE_k = \d_{jk}I \qiq I = \Sk E_k^TE_k $$ The $\{E_k\}$ basis can be used to construct $A$ as a partitioned matrix $$\eqalign{ A &= \Sj\Sk E_j\,a_{jk}\,E_k^T = \m{ a_{11}&a_{12}&\ldots&a_{1n} \\ a_{21}&a_{22}&\ldots&a_{2n} \\ \vdots&\vdots&\ddots&\vdots \\ a_{n1}&a_{n2}&\ldots&a_{nn} \\ } }$$ and to extract the individual partitions, e.g. $\;\;a_{jk} = E_j^T A E_k \,\in\, W$

The partial trace replaces each $\{a_{jk}\}$ block with its trace $$\eqalign{ P &= \ptrace{A} \:\in\: V \\ &= \Sj\Sk \e_j\,\e_k^T\,P_{jk} \\ &= \Sj\Sk \e_j\,\e_k^T\,\trace{a_{jk}} \\ &= \Sj\Sk \e_j\,\e_k^T\,\trace{E_j^T A E_k} \\ &= \Sj\Sk \e_j\,\e_k^T\,\LR{E_jE_k^T:A} \\ }$$ The gradient of the partial trace is a fourth-order tensor, which can be expressed in a variety of ways $$\eqalign{ \grad{P}{A} &= \Sj\Sk\LR{\e_j\,\e_k^T}&\star\LR{E_jE_k^T} \;\in\;\bbR{n\times n\times p\times p} \\ \grad{P}{a_{jk}} &= \LR{\e_j\,\e_k^T}\star I \;&\in\;\bbR{n\times n\times m\times m} \\ \grad{P_{jk}}{A} &= {E_jE_k^T} \;&\in\;\bbR{p\times p} \\ \grad{P_{jk}}{a_{rs}} &= \d_{jr}\d_{ks}\:I \;&\in\;\bbR{m\times m} \\ }$$ where $(\star)$ is the dyadic/tensor product and $(:)$ is the Frobenius product, i.e. $$\eqalign{ A:B &= \Sj\Sk A_{jk}B_{jk} \;=\; \trace{A^TB} \\ A:A &= \|A\|^2_F \qquad \{ {\rm Frobenius\:norm} \} \\ }$$ Define the matrix variable $$\eqalign{ &X = \a A+\b B \qiq dX = A\,d\a \\ &\LR{\a,\b \:\in\: \bbR{\tt1}} \\ }$$ Since $X\:($like $A$ and $B)$ is Hermitian, it can be diagonalized by a unitary matrix $$X=QZQ^H,\qquad Q^HQ=I,\quad Z=\Diag{z}\in\bbR{p\times p}$$ This allows us to invoke the Daleckii-Krein Theorem $$\eqalign{ &F = f(X) \qiq Q^H\,dF\,Q = R\odot\LR{Q^H\,dX\,Q} \\ &R_{jk} = \begin{cases} {\Large\frac{f(z_i)-f(z_j)}{z_i-z_j}} \qquad{\rm if}\;z_i\ne z_j \\ \\ \quad f'(z_j) \qquad\qquad {\rm otherwise} \\ \end{cases} }$$ where $(\odot)$ denotes the elementwise/Hadamard product.

$\sf NB\!:$ For the current problem $f(z)=\exp(-iz),\quad f'(z) = -i\,f(z)$

Identifying $U=F,\,$ we have everything we need to differentiate the objective $$\eqalign{ P &= \ptrace{UK} \\ &= \Sj\Sk\LR{\e_j\e_k}\star\BB{\BR{E_jE_k^T}:FK} \\ P_{jk} &= {E_jE_k^T}:FK \\ dP_{jk} &= {E_jE_k^T}:dF\,K \\ &= {E_jE_k^T}:Q\BR{R\odot\LR{Q^H\,\c{dX}\,Q}}Q^HK \\ &= {E_jE_k^T}:Q\BR{R\odot\LR{Q^H\c{A}Q}}Q^HK\:\c{d\a} \\ \grad{P_{jk}}{\a} &= {E_jE_k^T} : Q\BR{R\odot\LR{Q^HAQ}}Q^HK \\ }$$ which is the required gradient.