$$
\def\E{{\mathbb{E}}}
\def\R{{\mathbb{R}}}
\def\I{{\mathbb{I}}}
\def\P{{\mathcal{P}}}
\def\F{{\mathcal{F}}}
\def\rmd{{\mathrm{d}}}
\def\m{{{y}}}
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\def\U{{U}}
\def\u{{u}}
\def\rmd{{\mathrm{d}}}
$$

## TODO
- autoregressive horizon as an argument [Federico]
- metrics [Federico]
- MSE (on vorticity)
- Correlation
- approximate PDE residual [Emile]
- refactoring
- refactoring eval_score.py [Federico]
- Disantangling likelihood from guidance [Emile]
- Likelihood as a class
- Guidance as a class
- Move evaluation to main.py [Emile]
- load model's checkpoint
- subset of kolmogorov dataset for debugging purposes [Federico]
- neural network
- ModernUNet from PDERefiner [Emile]

- correction term between markov blanket score vs full score?
- momentum conservation by itself?

## Reconstruction guidance

### Setting

- We assume a likelihood $p(y|x_0)$, similar to the 'forward' operator in inverse problems.
- Generally, it could be given via a non linear mapping $y = A(x_0)$, although it is in many instances linear $y = A x_0$, as for inpainting, Gaussian blurring, etc
- the simplest is to assume a delta mass on this constraint $p(y|x_0) = \delta_{y=A(x_0)}$, although we can also 'soften' the constraint by adding some Gaussian noise on the observation, i.e. $p(y|x_0) = \mathrm{N}(y| A(x_0), \sigma_y^2)$

- We are interested in sampling $p(x_0|y)$, which can be done by sampling/simulating the following
- (backward) SDE $\rmd \overleftarrow{X}_t = \left\{ -f(\overleftarrow{X}_t) + \nabla{x_t} \log p_{T-t}(\overleftarrow{X}_t|y) \right\} \rmd t + g(T - t) \rmd B_t$,
- assuming a forward/noising SDE $\rmd X_t = f(X_t) \rmd t + g(t) \rmd B_t$

### Guidance term
- With Bayes rule, the _conditional score_ is given by $\nabla{x_t} \log p(x_t|y) = \nabla{x_t} \log p(x_t) + \nabla{x_t} \log p(y|x_t)$,
- where $\nabla{x_t} \log p(x_t)$ would be approximated by a score network $s_\theta(t, x_t) \approx \nabla{x_t} \log p(x_t)$ trained on the prior (with no information about the observation $y$).
- The conditioning therefore comes into play via $\nabla{x_t} \log p(y|x_t)$ referred as _reconstruction guidance_.
- How can we estimate p(y|x_t)? We have $p(A X_0 = y \vert x_t) = \int p(A X_0 = y \vert x_0) p_{0|t}(x_0|x_t) dx_t = \int \I_{A x_0 = \m}(x_0) p_{0|t}(x_0|x_t) dx_t$
but $p_{0|t}(x_0|x_t)$ is not available in closed-form. Yet, thanks to Tweedie's formula we have access to its moments. With $p_{t|0}(x_t|x_0) = \mathrm{N}(x_t|\mu_t x_0, \sigma_t^2 \mathrm{I})$, we have
- $\E[x_0|x_t] = \frac{x_t+\sigma_t^2 \nabla \log p_t(x_t)}{\mu_t} \triangleq \hat{x}(x_t)$
- $\mathrm{Cov}[x_0|x_t] = \left[\frac{\sigma_t^2}{\mu_t^2}(\mathrm{I} + \sigma_t^2 \nabla^2 \log p_t(x_t)) \right] \triangleq \hat{\Sigma}(x_t)$
- Hence motivating the Gaussian approximation $p(x_0|x_t) \approx q(x_0|x_t) \triangleq \mathrm{N}(x_0|\hat{x}(x_t), \hat{\Sigma}(x_t))$.
- Thus we can approximate the guidance term as
- $p(y \vert x_t) \approx q(y \vert x_t) \triangleq \int \mathrm{N}(y| A x_0, \Sigma_y) q(x_0|x_t) dx_0 = \mathrm{N}(y|A \hat{x}(x_t), A \hat{\Sigma}(x_t) A^\top + \Sigma_y)$, thus
- To approximate $\nabla_{x_t} \log p(y \vert x_t)$:
- Can get an _unbiased_ estimator of $\mathrm{det}(\hat{\Sigma}(x_t))$ [Sohl-Dickstein], yet _biased_ when taking the $\log$.
- Could compute $\hat{\Sigma}(x_t)^{-1} (y - A \hat{x}(x_t))$ via conjugate gradient, and use 'cheap' jvp inside.

### Choice of covariance estimation for $\mathrm{Cov}[x_0|x_t]$

- [Chung et al., 2022a]: $\mathrm{Cov}(y|x_t) \approx \texttt{detatch}(\|A(\hat{x}(x_t) - y\|) / \alpha$ with $\alpha$ a hyperparameter
- [Adam et al, 2022, Meng et al. 2022, Wu et al. 2023] rely on a diagonal covariance $\hat{\Sigma}(x_t) = \frac{\sigma_t^2}{\mu_t^2} \mathrm{I}$.
- [Finzi et al, 2023]:
- They assume $\Sigma_y = 0$ and $\sigma_t^2 \nabla^2 \log p_t(x_t)) \approx 0$
- Thus $\hat{\Sigma}(x_t) \approx (\sigma_t^2 / \mu_t^2)$
- Which they rescale with a hyperparameter $1/\alpha$
- Alternatively, approximating $\sigma_t^2 \nabla^2 \log p(x_t) \approx \lambda \mathrm{I}$, yields $\hat{\Sigma}(x_t)^{-1} \approx \frac{\mu_t^2}{(1+\lambda)\sigma_t^2} \mathrm{I}$
- [Rozet and Louppe] suggests using
- $\mathrm{Cov}(y|x_t) \approx \Sigma_y + \frac{\sigma_t^2}{\mu_t^2} A \Gamma A^\top$ since
- Assuming $p(x_0)$ to be Gaussian: $\hat{\Sigma}(x_t) \approx \Sigma_x - \Sigma_x (\Sigma_x + \frac{\sigma_t^2}{\mu_t^2} \mathrm{I})^{-1} \Sigma_x = \frac{\sigma_t^2}{\mu_t^2} Q \Lambda (\Lambda + \frac{\sigma_t^2}{\mu_t^2} \mathrm{I})^{-1} Q^{-1}$ with $Q \Lambda Q^{-1} = \mathrm{Cov}(x_0)$
- For most times $t$, $\Lambda (\Lambda + \frac{\sigma_t^2}{\mu_t^2} \mathrm{I})^{-1} \approx \Lambda (\Lambda + \mathrm{I})^{-1} \triangleq \Gamma$, thus $\hat{\Sigma}(x_t) \approx \frac{\sigma_t^2}{\mu_t^2} \Gamma$
- In practice $\Gamma = \gamma \mathrm{I}$ with $\gamma$ being an hyperparameter

### Exact score for observations
- As suggested in [Finzi et al, 2023] (in the codebase), when the forward operator $A$ is selecting a subset of $x$ (i.e 'slicing' or 'indexing' as in e.g. inpainting) then we know in closed form the score of the observed variables
- Let's define $A: x \mapsto x[I]$ with $I$ the indices of observed variables
- $\nabla_{x_t[I]} p(x_t|y) = \nabla_{y_t} p(y_t|y_0)$ with $y_t = x_t[I]$ and $y_0=y$ to stress the analogy with the conditional score appearing in the denoising score matching loss
- Then we can override $\nabla_{x_t} p(x_t|y)[I] = \nabla_{x_t[I]} p(x_t[I]|y)$ with the RHS being available in close form.

### Guidance to low PDE residuals

Let's assume a dynamical system $(\u^t)_{t \ge 0}$ with $\u^t \in \mathcal{X}$, potentially defined as the solution of an SPDE $\P(\u^t) = \mathcal{W}$, with $\P: \mathcal{X} \times \R \rightarrow \mathcal{X} \times \R$ a differential operator, and $\mathcal{W}$ white noise.

#### Divergence-free velocity (incompressible fluid)
We have $\mathrm{div}(\u^t) = 0$, which in Fourier space, is given by $\F^{-1}\left(i \xi \cdot (\F \u^t)(\xi)\right) = 0$

#### Navier-Stokes (incompressible fluid)
Let's denote the vorticity $\omega = \mathrm{curl} (\u)$.
Then we have the following PDE
$$ \P(\u) = \frac{\partial \omega}{\partial t} + u \cdot \nabla \omega - \frac{1}{Re} \Delta \omega - f = 0. $$
- The first term is approximated by finite difference: $\frac{\partial \omega}{\partial t} \approx \frac{\omega^{t+1} - \omega^{t-1}}{2 * \delta}$ with $\delta$ the step size.
- The next two terms are approximated in the Fourier space
- $\Delta \omega = \F^{-1}\left(-||\xi||^2 \F \omega \right)$
- $\u \cdot \nabla \omega = \u \cdot \F^{-1}\left(i \xi \ \F \omega \right)$

#### Noise
Assuming some white noise $\mathcal{W}$ with covariance $\Sigma_y$, we get that the SPDE induces the following guidance likelihood
$p(\u) = \mathcal{N}(0|\P(\u), \Sigma_y)$, in practice we assume $\Sigma_y = \sigma_y \mathrm{Id}$.

## Infinite dimensional exposition

Let's assume a dynamical system $(\U^t)_{t \ge 0}$ with $\U^t \in \mathrm{C}(\mathcal{X}, \mathcal{Y})$, potentially defined as the solution of an SPDE $\P(\U^t) = \mathcal{W}$, with $\P: \mathrm{C}(\mathcal{X} \times \R, \mathcal{Y}) \rightarrow \mathrm{C}(\mathcal{X} \times \R, \mathcal{Y})$ a differential operator, and $\mathcal{W}$ white noise.

- **Goal**: The goal is that, given some observations $(\U^{t_1}, \dots, \U^{t_m})$, one is able to make forecast / prediction $(\U^{t_{m+1}}, \dots, \U^{t_N})$

- **Training**: The idea is to train a diffusion model over the joint $(\U^t)_{t \ge 0}$, potentially on discretised time state $(\U^{t_1}, \dots, \U^{t_N})$.
Denoting by $\U_s$ samples paths from the SPDEs, either continuous, or discrete (in time and space), we define (accordingly) a noising diffusion process: $\rmd \U_s = -\frac{1}{2} \U_s g_s \rmd s + g^2_s \rmd \mathrm{B}_s$

- **Forecasting**:
Assuming we can fit $k$ states in memory, then we can autoregressively sample the conditionals $U^{t_{m+1}}|U^{t_{m}},\dots, U^{t_{m-k}}$

- **Conditioning on past trajectory**:
Past trajectory $(\dots, y^{t_{k-1}}, y^{t_k})$
- either amortise score network $s_\theta$
- or also condition at sampling time with $p(y|\u_0) = \I\left({(\dots, y^{t_{k-1}}, y^{t_k}) = (\dots, x^{t_{k-1}}_0, x^{t_k}_0)}\right) = \I(y = Ax_0)$, where $A$ is a linear operator masking the future time steps.
- $p(y|u_t) \approx \mathrm{N}(y|A\hat{u}(u_t), A~\hat{\Sigma}(u_t)~A^\top)$ by approximating $p(\u_0|\u_t) \approx \mathrm{N}(y|\hat{u}(u_t), \hat{\Sigma}(u_t))$ as Gaussian

- **Guiding towards PDE solutions**:
The idea to condition both on the trajectory and but also on the PDE residual, so as to 'guide' the solution to regions aproximately satisfying the PDE, and as such would prevent the solution to violate (too much) conserved quantities (e.g. energy).
- _PDE residual penalty_:
- $p(y|\u_0) \propto \exp\left(-\beta \|\P(\u_0)\|^2 \right) = \exp\left(-\beta \sum_{m} \|\P(\u_0^{t_m})\|^2 \right)$ yet cannot approximate $p(y|\u_t) = \int p(y|\u_0) p(\u_0|\u_t) \rmd x_t$ simply anymore??
- $p(y|\u_0) = \I(\P(\u_0) = 0)$, then approximate $p(y|u_t) \approx \mathrm{N}(y|\P(\hat{u}(u_t)), \nabla\P~\hat{\Sigma}(u_t)~\nabla\P^\top)$ by approximating $p(\u_0|\u_t) \approx \mathrm{N}(y|\hat{u}(u_t), \hat{\Sigma}(u_t))$ as a Gaussian and linearising $\P$ around $\hat{u}(u_t)$.
- _Computing PDE residual_:
- If $\U^t$ is modelled explicitely as a function (in space-time) (e.g. hyperpdiffusion), then we can exactly compute $\P(\u^t)$ via autodiff.
- If $\U^t$ is modelled continuously, yet only samples paths are available $u_0 \sim p$ (e.g. neural diffusion approaches), then can compute the differential operator by either differentiating through the solver (memory intensive), or by solving an adjoint ODE.
- If $\U^t$ is discretised, can approximate the time derivative with finite difference, and the spatial derivatives in Fourier space see [Shu et al., 2022].

- At sampling time the drfit of the condititonal process involved $\nabla{\u_t} \log p(x_t|y)$ which by Bayes rule is given by $\nabla{\u_t} \log p(x_t|y) = \nabla{\u_t} \log p(\u_t) + \nabla{\u_t} \log p(y|\u_t)$.
- Where $p(y|\u_t) = \int p(y|\u_0) p(\u_0|\u_t) \rmd x_t$ yet as $p(\u_0|\u_t)$ is unknown it needs to be approximxated see [Wu et al. 2023, Finzi et al, 2023,Adam et al, 2022, Meng et al. 2022, Wu et al. 2023, Rozet and Louppe].

### All-at-once vs autoregressive forecast
The fact that sampling all-at-once vs sampling in an autoregressive manner is giving different results is strange. I'll try to describe how I am sampling at the moment to see if there is some differences.

Let's assume we want to sample a 30-step long trajectory conditioned on the first 4 states, or observations from those states. Let's also assume that our pseudo Markov blanket is long $M$.

**All-at-once generation**We are interested in sampling from $p(x_{1:30}\mid y_{1:4})$ but our model was trained only to generate $p(x_{1:L})$. Using Bayes' rule we can write this as:
\begin{align}
p(x_{1:30}\mid y_{1:4}) = \frac{p(x_{1:30}, y_{1:4})}{p(y_{1:4})} = \frac{p( y_{1:4} \mid x_{1:30}) p(x_{1:30})}{p(y_{1:4})}
\end{align}

We can easily see that:
\begin{align}
\nabla_{x_{1:30}}\log p(x_{1:30}\mid y_{1:4}) = \nabla_{x_{1:30}} \log p( y_{1:4} \mid x_{1:30}) + \nabla_{x_{1:30}} \log p(x_{1:30})
\end{align}
where the first term is the *reconstruction guidance* and the second term is the score we trained our network to approximate.

As we can see from the equation above, the only guidance term we are using is given by $\nabla_{x_{1:30}} \log p( y_{1:4} \mid x_{1:30})$. Therefore, only the observations $y_{1:4}$ are affecting all the states. It would be interesting to understand how the information propagates to states that are further away from $x_{1:4}$.

<ins>Computational Complexity</ins>: the number of calls to the score networks depends on mostly two factors: 1) number of predictive steps (or time discretization) and 2) number of corrective Langevin samples. Assuming we are generating a single trajectory, then when evalauting the score network we unfold all the psuedo markov blanket and feed them in the same batch to the score network. So, assuming $D$ predictive steps and $C$ correction steps, then we have $D \times C$ calls to the score network.

<ins>Memory Complexity</ins>: since we are generating everyhting all together, we have to keep in memory the full trajectory, which corresponds to $\text{number of trajectory steps} \times \text{state dimension}$.

**Autoregressive generation** In sampling $p(x_{1:30}\mid y_{1:4})$ in an autoregressive manner, we are currently doing the following steps:

1. Step: in the example above, since we are considering a predictive horizon of only 1, we start by generating the first full psuedo Markov blanket, i.e. we are interested in
\begin{align}
\nabla_{x_{1:M}} \log p(x_{1:M}\mid y_{1:4}) = \nabla_{x_{1:M}} \log p(y_{1:4} \mid x_{1:M}) + \nabla_{x_{1:M}} \log p(x_{1:M})
\end{align}

2. After that, we keep repeating the following:
\begin{align}
\nabla_{x_{2:M+1}} \log p(x_{2:M+1}\mid y_{2:5}) = \nabla_{x_{2:M+1}} \log p(y_{2:5} \mid x_{2:M+1}) + \nabla_{x_{2:M+1}} \log p(x_{2:M+1})
\end{align}

3. At the next step we do
\begin{align}
\nabla_{x_{3:M+2}} \log p(x_{3:M+2}\mid y_{3:6}) = \nabla_{x_{3:M+2}} \log p(y_{3:6} \mid x_{3:M+2}) + \nabla_{x_{3:M+2}} \log p(x_{3:M+2})
\end{align}
5. and so on.. until we generate the full trajectory. In case of prediction horizon of 1, we are computing the following:
\begin{align}
\nabla_{x_{i:M+2}} \log p(x_{i:M+(i-1)}\mid y_{i:i+Q}) = \nabla_{x_{i:M+2}} \log p(y_{i:i+Q} \mid x_{i:M+(i-1)}) + \nabla_{x_{i:M+2}} \log p(x_{i:M+(i-1)})
\end{align}
where $Q$ is the number of conditioned frames.

Therefore, we can notice how while the prior score is almost similar to the prior score we are computing in the *all-at-once* setting, although here we are generate the same state several times depending on the prediction horizon. However, the reconstruction guidance in this setting is getting stronger, as it accounts for more terms that are not used in the all-at-once generation.

<ins>Computational Complexity</ins>: In this setting, at each step we have to call the score network $D \times C$ times, althoug since the batch is smaller, it tends to be a bit faster. However, we have to do this for $S$ steps, where the number of steps depends on the following parameters: 1) the prediction horizon, 2) the pseudo markov blanket size, and 3) the trajectory length.

<ins>Memory Complexity</ins>: in this case we have to just keep in memory $(\text{conditioned_frame} + \text{prediction horizon}) \times \text{state dimensionality}$, which in case $\text{conditioned_frame} + \text{prediction horizon} <= M$ then becomes $M \times \text{state dimensionality}$.

## LGSSM
**1D - Model**
\begin{align}
x_i = ax_{i-1} + b\epsilon_i, \epsilon_i \sim \mathcal{N}(0, 1)\\
y_i = cx_i + d\epsilon_i, \epsilon_i \sim \mathcal{N}(0, 1)
\end{align}

The marginals $p(x_i)$ are all Gaussians, hence we can find the steady-state mean and variance:

\begin{align}
\mathbb{E}[x_i] = a\mathbb{E}[x_{i-1}] + b\mathbf{\epsilon_i} \implies \mu = a\mu \implies \mu = 0
\end{align}

\begin{align}
var[x_i^2] = a^2var[x_{i-1}] + b^2 var[\epsilon_i] \implies \sigma^2 = a^2\sigma^2 + b^2 \implies \sigma^2 = \frac{b^2}{1-a^2}
\end{align}

\begin{align}
cov[x_i, x_{i-1}] = cov[ax^2_{i-1} + bx_{i-1}\epsilon_i] = acov[x^2_{i-1}] = a\sigma^2
\end{align}

In general, $cov[x_i, x_{i-k}] = a^{|k|}\sigma^2$

Hence, the joint distribution of the chain is a multivariate normal with the following mean and covariance matrix:

\begin{align}
p([x_1, x_2, ... x_N])
&= \mathcal{N}([0, 0, ..., 0], \begin{bmatrix}
\sigma^2 & a\sigma^2 & a^2\sigma^2 & \dots & a^{N-1}\sigma^2 \\
a\sigma^2 & \sigma^2 & a\sigma^2 & \dots & a^{N-2}\sigma^2 \\
\vdots & \vdots & \vdots & \ddots & \vdots \\
a^{N-1}\sigma^2 & a^{N-2}\sigma^2 & a^{N-3}\sigma^2 & \dots & \sigma^2
\end{bmatrix}) \\
&= \mathcal{N}([0, 0, ..., 0], \Sigma_0)
\end{align}

To then get to the diffused states, we have the following:

\begin{align}
x_{i-k:i+k}(t) = \mu(t)x_{i-k:i+k} + \sigma(t)\epsilon
\end{align}

with $\mu(t)=\alpha(t)$ and $\sigma(t)=\sqrt{1 - \alpha(t)^2}$, with $\alpha(t) = e^{-1/2 \int^t \beta(s) \mathrm{d}s}$ and $\beta$ the noising schedule.

$\EMILE$
$\E[x(t)] = \mu(t) \E[x(0)] = \alpha(t) \E[x(0)] = 0$
$\mathrm{Cov}[x(t)] = (1 - \sigma^2(t)) \mathrm{Cov}[x(0)] + \sigma^2(t) \mathrm{Id} = \alpha(t)^2 \mathrm{Cov}[x(0)] + (1 - \alpha^2(t)) \mathrm{Id}$

Leading to

$p([x_{i-k}(t), ... x_i(t), ... x_{i+k}(t)]) =$
\begin{align}
&=\mathcal{N}([0, 0, ..., 0], \alpha(t)^2 \Sigma_0 + (1 - \alpha^2(t)) \mathrm{Id} ) \\
&=\mathcal{N}([0, 0, ..., 0], \mu(t)^2 \Sigma_0 + \sigma^2(t) \mathrm{Id} ) \\
&=\mathcal{N}([0, 0, ..., 0], \begin{bmatrix}
\mu(t)^2\sigma^2 + \sigma^2(t) & \mu(t)^2a\sigma^2 & \mu(t)^2a^2\sigma^2 & \dots & \mu(t)^2a^{2k}\sigma^2 \\
\mu(t)^2a\sigma^2 & \mu(t)^2\sigma^2 + \sigma^2(t) & \mu(t)^2a\sigma^2 & \dots & \mu(t)^2a^{2k-1}\sigma^2 \\
\vdots & \vdots & \vdots & \ddots & \vdots \\
\mu(t)^2a^{2k}\sigma^2 & \mu(t)^2a^{2k-1}\sigma^2 & \mu(t)^2a^{2k-3}\sigma^2 & \dots & \mu(t)^2\sigma^2 + \sigma^2(t)
\end{bmatrix})
\end{align}

**All-at-once sampling** For all-at-once sampling we need to find $\nabla_{x_i(t)}log(p(x_{i-k:i+k}(t)))$ (because this is what we are training our score network on $s_{\phi}(x_{i-k:i+k}(t), t) \approx \nabla_{x_i(t)}log(p(x_{i-k:i+k}(t)))$.

$\nabla_{x_i(t)}log(p(x_{i-k:i+k}(t))) =$

\begin{align}
=\nabla_{x_i(t)}\bigg[-\frac{1}{2}\begin{bmatrix}
x_{i-k}(t) \\
\vdots \\
x_{i+k}(t)
\end{bmatrix}^T\begin{bmatrix}
\mu(t)^2\sigma^2 + \sigma^2(t) & \mu(t)^2a\sigma^2 & \mu(t)^2a^2\sigma^2 & \dots & \mu(t)^2a^{2k}\sigma^2 \\
\mu(t)^2a\sigma^2 & \mu(t)^2\sigma^2 + \sigma^2(t) & \mu(t)^2a\sigma^2 & \dots & \mu(t)^2a^{2k-1}\sigma^2 \\
\vdots & \vdots & \vdots & \ddots & \vdots \\
\mu(t)^2a^{2k}\sigma^2 & \mu(t)^2a^{2k-1}\sigma^2 & \mu(t)^2a^{2k-3}\sigma^2 & \dots & \mu(t)^2\sigma^2 + \sigma^2(t)
\end{bmatrix}^{-1}\begin{bmatrix}
x_{i-k}(t) \\
\vdots \\
x_{i+k}(t)
\end{bmatrix}\bigg]
\end{align}

\begin{align}
=\nabla_{x_i(t)}\bigg[-\frac{1}{2}\begin{bmatrix}
x_{i-k}(t) \\
\vdots \\
x_{i+k}(t)
\end{bmatrix}^T\mathbf{A}\begin{bmatrix}
x_{i-k}(t) \\
\vdots \\
x_{i+k}(t)
\end{bmatrix}\bigg]
\end{align}

We can use Einstein notation to get:

\begin{align}
\nabla_{x_i(t)}[-\frac{1}{2}\sum_{p=-k}^{k}\sum_{l=-k}^{k}x_{i-p}A_{k+1-p, k+1-l}x_{i-l}] &= -\frac{1}{2}[\sum_{l=-k}^{k}A_{k+1, k+1-l}x_{i-l} + \sum_{p=-k}^{k}A_{k+1-p, k+1}x_{i-p}]\\
& = -\frac{1}{2}[\sum_{p=-k}^{k}(A_{k+1, k+1-p} + A_{k+1-p, k+1})x_{i-p}]
\end{align}

In this case we only obtain the gradient wrt middle of the blanket, but the general formula wrt to any element of the blanket is:
\begin{align}
\nabla_{x_i+s(t)}[-\frac{1}{2}\sum_{p=-k}^{k}\sum_{l=-k}^{k}x_{i-p}A_{k+1-p, k+1-l}x_{i-l}] &= -\frac{1}{2}[\sum_{l=-k}^{k}A_{k+s+1, k+1-l}x_{i-l} + \sum_{p=-k}^{k}A_{k+1-p, k+s+1}x_{i-p}]\\
& = -\frac{1}{2}[\sum_{p=-k}^{k}(A_{k+s+1, k+1-p} + A_{k+1-p, k+s+1})x_{i-p}]
\end{align}

**Implementation** In the actual implementation the blanket looks like $x_0(t), x_1(t), ..., x_{2k}(t)$, so the corresponding formulas are:

\begin{align}
\nabla_{x_i(t)}[-\frac{1}{2}\sum_{p=0}^{2k}\sum_{m=0}^{2k}x_{p}A_{p, m}x_{m}] = -\frac{1}{2}[\sum_{m=0}^{2k}A_{i,m}x_m + \sum_{p=0}^{2k}A_{p, i}x_p] = -\frac{1}{2}[\sum_{p=0}^{2k}(A_{p, i} + A_{i, p})x_p]
\end{align}

**Auto-regressive sampling** For example, for the AR sampling, we will need to compute $\nabla_{x_{i+1}(t)}log(p(x_{i-k:i+k}(t)))$.

**Multidimensional - Model**
$\mathbf{X}_i = \mathbf{T}\mathbf{X}_{i-1} + \mathbf{\epsilon_T}$, where $\mathbf{\epsilon_T}\sim\mathcal{N}(\mathbf{0}, \mathbf{\Sigma_T})$

$\mathbf{Y}_i = \mathbf{O}\mathbf{X}_{i} + \mathbf{\epsilon_O}$, where $\mathbf{\epsilon_O}\sim\mathcal{N}(\mathbf{0}, \mathbf{\Sigma_O})$

The marginals $p(\mathbf{X}_i)$ are all multivariate Gaussians, with mean and covariance:

$\mathbb{E}[\mathbf{X}_i] = \mathbf{T}\mathbb{E}[\mathbf{X}_{i-1}] + \mathbb{E}[\mathbf{\Sigma_T}] \implies \mathbf{\mu} = \mathbf{T}\mathbf{\mu} \implies \mathbf{\mu} = \mathbf{0}$

$var[\mathbf{X}_i] =\mathbf{T}var[\mathbf{X}_{i-1}]\mathbf{T}^T + \mathbf{\Sigma_T} \implies vec(var[\mathbf{X}_i]) = (\mathbf{I} - T\bigotimes T)^{-1} - vec(\mathbf{\Sigma_T})$
where for $\mathbf{A}$ matrix of $m\times n$, $vec(\mathbb{A}) = [a_{11}, ...a_{m1}, a_{12}, ... a_{m2}, ... a_{mn}]$

Let the covariance be $\Sigma_{SS}$, then
$p(\mathbf{X}_i) = \mathcal{N}(\mathbf{0}, \Sigma_{SS})$

The covariance between consecutive terms:

$cov[\mathbf{X}_i, \mathbf{X}_{i-1}] = cov[\mathbf{T}\mathbf{X}_{i-1}\mathbf{X}_{i-1}^T + \mathbf{X}_{i-1}\mathbf{\epsilon_T}^T]$, so $cov[\mathbf{X}_i, \mathbf{X}_{i-1}] = \mathbf{T}\Sigma_{SS}$

And $cov[\mathbf{X}_i, \mathbf{X}_{i-k}] = \mathbf{T}^k\Sigma_{SS}$

## Prior work
- [Rozet et al, Finzi et al., Shu et al., Ho et al.]
- $\TODO$ organise prior work with labels / key words

## References
- Alexandre Adam et al. “Posterior samples of source galaxies in strong gravitational lenses with score-based priors”. 2022. URL: https://arxiv.org/abs/2211.03812.
- Xiangming Meng and Yoshiyuki Kabashima. “Diffusion Model Based Posterior Sampling for Noisy Linear Inverse Problems”. 2022. URL: https://arxiv.org/abs/2211.12343.
- Finzi, Marc Anton and Boral, Anudhyan and Wilson, Andrew Gordon and Sha, Fei and {Zepeda-Nunez}, Leonardo. User-Defined {{Event Sampling}} and {{Uncertainty Quantification}} in {{Diffusion Models}} for {{Physical Dynamical Systems}}. 2023. URL: https://openreview.net/forum?id=sdhcjMzhHN.
-Rozet, Fran{\c c}ois and Louppe, Gilles. Score-Based {{Data Assimilation}}. 2023. URL: http://arxiv.org/abs/2306.10574
- Wu, Luhuan and Trippe, Brian L and Naesseth, Christian A and Cunningham, John P. Practical and {{Asymptotically Exact Conditional Sampling}} in {{Diffusion Models}}. 2023.
- Jascha Sohl-Dickstein. Two Equalities Expressing the Determinant of a Matrix in Terms of Expectations over Matrix-Vector Products. 2020. URL: http://arxiv.org/abs/2005.06553.
- Mathieu et al. Geometric Neural Diffusion Processes. 2023 https://arxiv.org/abs/2307.05431
- Dule Shu, Zijie Li, Amir Barati Farimani, A Physics-informed Diffusion Model for High-fidelity Flow Field Reconstruction
- Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, David J. Fleet. Video Diffusion Models
- Hyungjin Chung, Jeongsol Kim, Michael T. Mccann, Marc L. Klasky, Jong Chul Ye, Diffusion Posterior Sampling for General Noisy Inverse Problems
- Chung, H., Sim, B., Ryu, D., and Ye, J. C., Improving diffusion models for inverse problems using manifold constraints., 2022