How to calculate the gradient and Hessian for composite function? I have a function that can be written as follows:
$$f(\vec{x},\vec{y})  = f_x(\vec{x}) + f_y(\vec{y}) + k\bigg(f_x(\vec{x})-f_y(\vec{y})\bigg)^2$$
I have to take the first derivative (gradient) and second derivative (hessian) of this function with respect to the combined $\vec{x}$ and $\vec{y}$. By combined, I mean a vector which the concatenation of both $\vec{x}$ and $\vec{y}$ like this:
$$\vec{z} = (x_1,x_2,x_3,...,y_1,y_2,y_3,...)$$
So far I have been able to derive the analytic expression for the gradient $\nabla f$, which is simple because $f_x$ depends only on $x$ terms, and $f_y$ depends only on y terms. So, I can simply expand the square expression, collect the gradient terms, and concatenate. However, I am struggling to derive an expression for the Hessian matrix ($\nabla^2 f$) because of the cross-terms.
Note: The gradient and Hessian of each of the $f_x(\vec{x})$ and $f_y(\vec{y})$ with respect to their own vectors are possible to calculate (available). So, I have to write the gradient and Hessian of the composite function in those terms, as I need those for a program I am coding. Both $f_x$ and $f_y$ are functions that produce scalar numbers as ouptut. $k$ is a scalar constant.
I am a chemist, so I am not very familiar with linear algebra. If the Hessian (or gradient) can be calculated quickly with a matrix operation then it would be easier to code for. Any help is appreciated.
 A: Let $h(x,y) = f_x(\vec{x})+f_y(\vec{y})+k\Big(f_x(\vec{x})-f_y(\vec{y})\Big)^2$.
Then
$$\nabla_{\vec{x},\vec{y}}h(\vec{x},\vec{y})=\begin{bmatrix}\nabla_\vec{x}\\\nabla_\vec{y}\end{bmatrix}h(\vec{x},\vec{y})\\
=
\begin{bmatrix}\nabla_\vec{x}f_x(\vec{x}) + k\nabla_\vec{x}f_x(\vec{x})^2-2kf_y(\vec{y})\nabla_\vec{x}f_x(\vec{x})
\\
\nabla_\vec{y}f_y(\vec{y}) + k\nabla_\vec{y}f_y(\vec{y})^2-2kf_x(\vec{x})\nabla_\vec{y}f_y(\vec{y})\end{bmatrix}$$
and
$$H_{\vec{x},\vec{y}}h(\vec{x},\vec{y})=\begin{bmatrix}\partial^2_{\vec{x},\vec{x}} & \partial^2_{\vec{x},\vec{y}} \\
\partial^2_{\vec{y},\vec{x}} & \partial^2_{\vec{y},\vec{y}}
\end{bmatrix}h(\vec{x},\vec{y})$$
$$=\begin{bmatrix}
H_\vec{x}f_x(\vec{x}) + kH_\vec{x}f_x(\vec{x})^2-2kf_y(\vec{y})H_\vec{x}f_x(\vec{x}) & -2k\nabla f_x^T \nabla f_y \\
-2k\nabla f_y^T \nabla f_x & H_\vec{y}f_y(\vec{y}) + kH_\vec{y}f_y(\vec{y})^2-2kf_x(\vec{x})H_\vec{y}f_y(\vec{y}) 
\end{bmatrix}$$
I've written these in block matrix notation. If you can compute the regular gradients and hessians of $f_i$ and $f_i^2$, then these formulae should get you to the gradient and hessian of $h$
A: $
\def\BR#1{\Big(#1\Big)}
\def\LR#1{\left(#1\right)}
\def\op#1{\operatorname{#1}}
\def\trace#1{\op{Tr}\LR{#1}}
\def\qiq{\quad\implies\quad}
\def\qif{\quad\iff\quad}
\def\a{\alpha}\def\b{\beta}\def\T{{\large\tt0}}
\def\l{\lambda}\def\t{\lambda}\def\p{\partial}
\def\grad#1#2{\frac{\p #1}{\p #2}}
\def\hess#1#2{\frac{\p^2 #1}{\p #2^2}}
\def\m#1{\left[\begin{array}{r}#1\end{array}\right]}
\def\mr#1{\left[\begin{array}{l}#1\end{array}\right]}
\def\mm#1{\left[\begin{array}{c|c}#1\end{array}\right]}
\def\c#1{\color{red}{#1}}
\def\CLR#1{\c{\LR{#1}}}
\def\fracLR#1#2{\LR{\frac{#1}{#2}}}
\def\o{{\tt1}}
\def\bbR#1{{\mathbb R}^{#1}}
$The inverse of concatenating matrix and vector variables is to partition them into blocks.
Towards that end, define block-matrix analogs of the cartesian basis vectors $\LR{e_k\in\bbR 2}$
$$\eqalign{
&E_1 = \mr{I_n\\ \T^T},\quad &E_2 = \mr{\T\\I_m}
 &\qiq &z = E_1x + E_2y = \m{x\\y} \\
&E_1^TE_1 = I_n,&E_1^TE_2 = \T &\qiq &x=E_1^Tz \\
&E_2^TE_1 = \T^T,&E_2^TE_2 = I_m &\qiq &y=E_2^Tz \\
&E_1E_1^T + E_2E_2^T &= I_{(m+n)} \\
&\T \in\bbR{n\times m} &x\in\bbR{n\times\o}
 \qquad &y\in\bbR{m\times\o}\qquad &z\in\bbR{(m+n)\times\o} \\
}$$
and give the known function/gradient/Hessian variables single-letter names
$$\eqalign{
\a &= \a(x),\qquad &a = a(x) = \grad{\a}{x},\qquad &A = A(x) = \grad{a}{x} = \hess{\a}{x} \\
\b &= \b(y),\qquad &b = b(y) = \grad{\b}{y},\qquad &B = B(y) = \grad{b}{y} = \hess{\b}{y} \\
}$$
Rewrite the function and calculate its gradient
$$\eqalign{
f &= \a + \b + k\LR{\a-\b}^2 \\
df &= d\a + d\b + \c{2k\LR{\a-\b}}\LR{d\a-d\b} \\
 &= d\a + d\b + \c{\t}\LR{d\a-d\b} \\
 &= \LR{\o+\t}\c{d\a} + \LR{\o-\t}\c{d\b} \\
 &= \LR{\o+\t}\c{a:dx} + \LR{\o-\t}\c{b:dy} \\
 &= \LR{\o+\t}a:E_1^Tdz + \LR{\o-\t}b:E_2^Tdz \\
 &= \LR{\o+\t}E_1a:dz + \LR{\o-\t}E_2b:dz \\
 &= \BR{\LR{\o+\t}E_1a + \LR{\o-\t}E_2b\,}:dz \\
\grad{f}{z}
 &= \BR{\LR{\o+\t}E_1a + \LR{\o-\t}E_2b\,}
 \;=\; g\;\Big({\rm gradient}\Big) \\
g &= \mm{\LR{\o+\t}a \\ \LR{\o-\t}b} \\
}$$
Before tackling the Hessian, here's the gradient of the $\l$ variable introduced above
$$\eqalign{
c &= \LR{E_1a - E_2b} = \m{a\\-b} \\
\t &= 2k\LR{\a-\b} \\
d\t &= 2k\LR{d\a-d\b} \\
 &= 2k\LR{E_1a-E_2b}:dz \\
 &= 2kc^Tdz \\
}$$
The Hessian is simply the gradient of the gradient, therefore
$$\eqalign{
dg
 &= \LR{\o+\t}E_1\,\c{da} + \LR{\o-\t}E_2\,\c{db}
  + \LR{E_1a-E_2b}\c{d\t} \\
  &= \LR{\o+\t}E_1\,\c{A\,dx} + \LR{\o-\t}E_2\,\c{B\,dy}
  + \LR{E_1a-E_2b}\c{2kc^Tdz} \\
  &= \LR{\o+\t}E_1AE_1^Tdz + \LR{\o-\t}E_2BE_2^Tdz
  + 2kcc^Tdz \\
  &= \BR{\LR{\o+\t}E_1AE_1^T
   + \LR{\o-\t}E_2BE_2^T
   + 2kcc^T}\,dz \\
\grad{g}{z}
  &= \BR{\LR{\o+\t}E_1AE_1^T + \LR{\o-\t}E_2BE_2^T + 2kcc^T}
 \;=\; H\;\Big({\rm Hessian}\Big) \\
H &= \mm{
2kaa^T+\LR{\o+\t}A & -2kab^T \\\hline
-2kba^T & 2kbb^T+\LR{\o-\t}B \\
} \;=\; \hess{\!f}{z} \\
\\
}$$

In the above $(:)$ denotes the Frobenius product, which is extremely useful in Matrix Calculus
$$\eqalign{
A:B &= \sum_{i=1}^m\sum_{j=1}^n A_{ij}B_{ij} \;=\; \trace{A^TB} \\
A:A &= \|A\|^2_{\c F}\qquad \Big({\rm \c{Frobenius}\:norm}\Big) \\
}$$
This is also called the double-dot or double contraction product.
When applied to vectors $(n={\tt1})$ it reduces to the standard dot product.
The properties of the underlying trace function allow the terms in a
Frobenius product to be rearranged in many interesting ways, e.g.
$$\eqalign{
A:B &= B:A \\
A:B &= A^T:B^T \\
C:\LR{AB} &= \LR{CB^T}:A &= \LR{A^TC}:B \\
}$$
A: For the cross-term of the Hessian, you can write
\begin{eqnarray}
\mathbf{g}_x
\equiv
\frac{\partial f}{\partial \mathbf{x}}
&=&
-2 k  f_y(\mathbf{y}) 
\frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}} \\
d\mathbf{g}_x
&=&
-2 k  f_y(\mathbf{y}) 
\frac{\partial^2 f_x(\mathbf{x})}{\partial \mathbf{x}^2}
d\mathbf{x}
-2 k  
\frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}}
\left(
\frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}}
\right)^T
d\mathbf{y}
\end{eqnarray}
and similarly
\begin{eqnarray}
\mathbf{g}_y
\equiv
\frac{\partial f}{\partial \mathbf{y}}
&=&
-2 k  f_x(\mathbf{x}) 
\frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}} \\
d\mathbf{g}_y
&=&
-2 k  f_x(\mathbf{x}) 
\frac{\partial^2 f_y(\mathbf{y})}{\partial \mathbf{y}^2}
d\mathbf{y}
-2 k  
\frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}}
\left(
\frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}}
\right)^T
d\mathbf{x}
\end{eqnarray}
The Hessian is easily found by regrouping terms
\begin{equation}
-2k
\begin{bmatrix}
f_y(\mathbf{y}) 
\frac{\partial^2 f_x(\mathbf{x})}{\partial \mathbf{x}^2}
&
\frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}}
\left(
\frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}}
\right)^T \\
\frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}}
\left(
\frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}}
\right)^T
&
f_x(\mathbf{x}) 
\frac{\partial^2 f_y(\mathbf{y})}{\partial \mathbf{y}^2}
\end{bmatrix}
\end{equation}
