---
$\usepackage{mathtools}$
$\usepackage[utf8]{inputenc}$
---
$\newcommand{\BK}[1]{ {\left( #1 \right)} }
\newcommand{\sqBK}[1]{ {\left[ #1 \right]} }
\newcommand{\curBK}[1]{ {\left\{ #1 \right\}} }
\newcommand{\bra}[1]{\langle #1 \rangle}
\newcommand{\p}[2]{\frac{\partial #1 }{\partial #2 }}
\newcommand{\d}[2]{\frac{d #1 }{d #2 }}
\newcommand{\normal}[2]{\mathcal{N}( #1 , #2 )}
\newcommand{\R}{\mathbb{R}}
\newcommand{\M}{\mathcal{M}}
\newcommand{\E}{\mathbb{E}}
\newcommand{\N}{\mathcal{N}}
\newcommand{\Q}{\mathbb{Q}}
\newcommand{\s}{\sigma}
\newcommand{\T}{\text{T}}
\newcommand{\dFdw}{\frac{\partial F}{\partial w}}
\newcommand{\dFdu}{\frac{\partial F}{\partial u}}
\newcommand{\dwdu}{\frac{\partial w}{\partial u}}
\newcommand{\ddFdw}{\frac{\partial^2 F}{\partial w^2}}
\newcommand{\ddFdu}{\frac{\partial^2 F}{\partial u^2}}
\newcommand{\ddwdu}{\frac{\partial^2 w}{\partial u^2}}
\newcommand{\m}{\mathbf}$

# Inverse Problems with Constrained Hamiltonian Monte Carlo (HMC)

Our goal is to infer some unknown quantity given some observations provided by a known model. The inverse problem can be understood as performing inference on the unnormalized posterior. Alternatively, an equivalent interpretation is to infer from the prior with constraints that account for data or observations. Here, inference on a constrained prior is achieved using a constrained HMC algorithm.

## Defining The Manifold

We assume a data model with additive noise,

$$y = \bar y + \sigma n$$

$$ \bar y = F(z)$$

where
$z \in \R^d$ is the unknown variable
$F: \R^d \rightarrow \R^m$ is the forward PDE operator on $z$
$n \sim \N(0, I_{m \times m}) \in \R^m$ is the (unknown) noise variable
$\s \in \R$ is noise intensity
$y \in\R^m$ is the vector of observation

This model take in an input of $d$ dimensions and gives us $m$ observations. Note that typically $d \gg m$, which gives rise to the charateristic ill-posedness of inverse problems. Taking the $m$ observations to be constraints, we can define the solution space as a manifold, $\M$ embedded in ambient space, $\R^{m+d}$.

$$ \M = \{(z,n)\} = \Big\{ \Big(z, \frac{y - F(z)}{\s} \Big) \Big\}$$

Special case: when $G$ is linear, $\M$ is a hyperplane in $\R^{d+m}$.

### Forward Operator, $F$

We consider a non-linear Poisson's equation with known source term $f$, observations of solution field $s$ and unknown log coefficient field $c(z)$ with Matérn covariance.

$$ \nabla \cdot ( e^{c} \nabla s ) = f $$

The weak form of the PDE is well known. Coloured noise $c$ is acquired by solving another SPDE. (here the SPDE is not stochastic but we call it SPDE to avoid confusion with the Poisson PDE)

$$ ( \kappa^2 - \Delta)^{\alpha/2} c(z) = w(z) $$

where scaling parameter $\kappa$ is arbitrarily chosen to be 0.1. For $\alpha = 2$, the weak form can also be easily written down. The white noise is given by $L^T w = z$, where $L$ is the cholesky factor of the mass matrix (of the finite element basis). We sample $z$ from the standard normal distribution, $\N(0,I_{d \times d})$.

Once we solve the Poisson's PDE, we use an observation operator to 'read' the solution. We approximate a dirac delta function with a Gaussian function with small variance.

$$ \bar y(x') = \int \delta(x-x') \, s(x) \, dx $$

where $x \in \R^2$ in our experiments.

Overall, the forward operator is a composition of the observation operator, Poisson PDE, the SPDE, and white noise process. The following diagram will be important later to intuit the gradient & hessian computations.

$$ z \rightarrow w \rightarrow c \rightarrow s \rightarrow \bar y$$

### Jacobian & Gram Matrix

Our constraint function is

$$ C(z,n) = F(z) + \s n - \bar y $$

We write down the Jacobian $J$ as it is ubiquitous throughout the constrained HMC algorithm.

$$ J = \Big[ \p{C}{z} \p{C}{n} \Big] = \Big[ \d{F}{z} \quad \sigma I \Big] $$

The gram matrix is given by

\begin{align}
J J^\T =\Big[ \d{F}{z} \d{F}{z}^\T + \sigma^2 I\Big] \in \R^{d \times d}.
\end{align}

Given $J$ is full rank, the inverse of the gram matrix always exists because it is positive semidefinite since $x^T(J J^\T)x = (Jx)^\T Jx \geq 0$.

We take the cholesky decompositon of the gram matrix, $LL^\T = JJ^\T$. This is convenient for operatios on the gram matrix as well as computing its determinant.

## Hamiltonian Dynamics

We put our state and noise variables in a joint vector $u = [ z \quad n ]$ on which we wish to perform inference. The auxiliary momentum variable is given as $p$.

To do HMC, we need:
1. The potential energy function, $\psi(u)$ and its gradient $\d{\psi}{u}$.
2. The kinetic energy function, $K(p)$ and its gradient $\d{K}{p}$.

The kinetic energy function used here is the standard $K(p) = \frac{p^\T p}{2}$. The gradient is straightfoward, $\d{K}{p} = p$. The potential energy function $\psi$ corresponds to the negative log posterior (target) density.

### Target density

Let the prior density on $u$ be $\pi(u)$, which is defined over $\R^{d}$. The corresponding density defined on $\M$ is then

$$ \pi_M = \pi |J^\T J|^{-1/2} .$$

### $\psi$ and $\d{\psi}{u}$

Taking the prior on $u$ to be a standard normal, the negative log posterior is given by

\begin{align}
\psi(u) &= -\text{log} \, \pi_M \\
&= -\text{log }\pi + \frac{1}{2}\text{log } |J J^\T| \\
&= Z + \frac{u^\T u}{2} + \sum_i \text{log } L_{ii}
\end{align}

Note that $Z$ is the normalising constant that is independent of $u$. Hence, it will disappear in gradient computations as well as cancel out in metropolis-hastings. So, we can forget about $Z$.

The gradient works to out be as such (I have not derived this myself).

$$\d{\psi}{u} = \d{}{z} \frac{1}{2} \text{log }\Bigg| \d{F}{z} \d{F}{z}^\T + \sigma^2 I\Bigg| = \text{Trace} \Bigg( \Big( \d{F}{z} \d{F}{z}^\T + \sigma^2 I\Big)^{-1} \d{F}{z} \frac{\partial^2 F^\T}{\partial z_i \partial z} \Bigg) $$

where the second derivative of $F$ with respect to $z$ is called the hessian. Note that $\psi$ is constnat with respect to noise variable, $n$.

## Constrained HMC Algorithm

The algorithm I used is from [Matt's paper](https://arxiv.org/pdf/1605.07826.pdf).

### Working with the Cholesky Decomposition

:::warning
Note to reader: The bold typeface (up to Computational Challenges) is a relic of copying from a LaTeX file, nonetheless the symbol conventions are consistent.
:::

Cholesky decomposition of mass matrix, $\m{C}$, gives

$$\m{L_A L_A^T = A = P C P^T}.$$

Every subsequent solve functions using the `sksparse.cholmod` library, , will be done with respect to matrix $\m{A}$. We rearrange to get $\m{(P^{-1} L_A) (P^{-1} L_A)^T = C}$. Hence, we can think of $\m{L_C = (P^{-1} L_A)}$ to be the decomposition for matrix $\m{C}$.

We are interested to solve for white noise via $\m{L_C^T w = u}$, or $\m{L_A ^T P^{-T} w = u}$. The function call `Factor.solve_Lt(b)` returns $\m{x}$ for $\m{L_A^T x=b}$. By calling the function above with $\m{u}$ as input, we get $\m{P^{-T} w}$, which is the input for `Factor.apply_Pt(b)` where we acquire $\m{w}$ by (as the function name suggests) applying $\m{P^T}$ in front.

### Computing the gradient

Let $\m{F(w(u))}$ be a forward operator for a PDE. For gradient computations,

$$\m{\dFdu = \dFdw \dwdu},$$

where $\m{\dwdu = P^T L_A^{-T}}$ since $\m{w = P^T L_A^{-T} u}$. (Simple proof: Let $\m{P^T L_A^{-T} = Q}$, then $\m{w_i = \sum_j Q_{ij} u_j}$. Therefore, $\m{\frac{\partial w_i}{\partial u_j} = Q_{ij}}$, from which we can see that the full derivative matrix is simply $\m{Q}$ itself.)\\

We can compute $\m{\dFdw}$ using the `compute_gradient()` function from `fenics_adjoint`. The rest of the computation works to be

\begin{align*}
\m{\dFdu} &= \m{\dFdw P^T L_A^{-T}} \\
&= \m{(\dFdw^T)^TP^T L_A^{-T} }\\
&= \m{(L_A^{-1} P \dFdw^T ) ^T} \\
\m{\dFdu^T} &= \m{(L_A^{-1} P \dFdw^T )}. \\
\end{align*}

To compute the right hand side of the final equation, we may first apply $\m{P}$ to $\m{\dFdw^T}$, then solve $\m{L_A x = P \dFdw^T}$ for $\m{x}$. That gives us the gradient $\m{\dFdu^T}$.

### Computing the hessian
Using chain rule, the hessian of $F$ with respect to $u$ works out to be

$$\m{H = \ddFdu = \frac{\partial}{\partial u} (\dFdu) = \frac{\partial}{\partial u} (\dFdw \dwdu) = \dFdw \ddwdu + \dwdu \frac{\partial F}{\partial w \partial u} = \dwdu^T \ddFdw \dwdu.}$$

Recall that $\m{w}$ is linearly related to $\m{u}$ so $\m{\ddwdu} = 0$. In reality, we do not need to compute the full hessian but only the action of the hessian on a vector $\m{v}$, i.e. a hessian-vector product,

$$\m{H v = \dwdu^T \ddFdw \dwdu v = L_A^{-1} P \ddFdw P^T L_A^{-T} v}.$$

Be reminded that the order of operations with which a matrix-matrix-vector multiplication is computed changes the overall complexity. Say to compute $\m{A B x }$, $\m{A (B x) }$ is preferable to $\m{ (A B) x }$ as the latter involves matrix-matrix multiplication, an $\mathcal{O}(n^3)$ operation. In our case, we ought to compute $\m{H v}$ from right to left, each time performing only a matrix-vector multiplication, an $\mathcal{O}(n^2)$ operation. Particularly, the action of the second derivative $\m{\ddFdw }$ on a vector can be computed via `compute_hessian()`. Be careful not to confuse the hessian in the function name with the hessian that we are interested in, the hessian in the function name refers to $\m{\ddFdw }$, where we later tag on the necessarily $\m{L_A^{-1}, P}$ terms (and their transposes) to account for linear factors of $\m{\dwdu}$.

### Computational Challenges

Recall the number of observations to be $m$.

At every $u$ update, the Jacobian computation calls `compute_gradient()` $m$ times.

At every $p$ update, the potential gradient computation calls `compute_gradient()` and `compute_hessian()` $m$ times each.

Loosely speaking, given $N$ MCMC samples and $N_s$ steps in each sample trajectory, there are about $2 \times N \times N_s \times m$ `compute_gradient()` calls and $N \times N_s \times m$ calls to `compute_hessian()`.

E.g. Assuming a problem of about 200 variables, which takes `compute_gradient()` about 0.04 seconds to run. Take 1000 MCMC samples, 5 steps per trajectory and 10 observations, that will amount to $10^5$ gradient calls, which amounts to a little over an hour just on gradient computations.

Reducing these numbers by some small factor is possible by storing the Jacobians and gradients at some avenues.

----

###### tags: `research` `MCMC`